[ { "query": "repo:ml-explore/mlx \"clip_grad_norm\"", "url": "https://api.github.com/search/issues?q=repo%3Aml-explore%2Fmlx+%22clip_grad_norm%22&per_page=100", "result": { "total_count": 7, "incomplete_results": false, "items": [ { "url": "https://api.github.com/repos/ml-explore/mlx/issues/3090", "repository_url": "https://api.github.com/repos/ml-explore/mlx", "labels_url": "https://api.github.com/repos/ml-explore/mlx/issues/3090/labels{/name}", "comments_url": "https://api.github.com/repos/ml-explore/mlx/issues/3090/comments", "events_url": "https://api.github.com/repos/ml-explore/mlx/issues/3090/events", "html_url": "https://github.com/ml-explore/mlx/pull/3090", "id": 3886885010, "node_id": "PR_kwDOKzRn187A9spW", "number": 3090, "title": "[WIP] Fuse into reduction", "user": { "login": "nastya236", "id": 41966024, "node_id": "MDQ6VXNlcjQxOTY2MDI0", "avatar_url": "https://avatars.githubusercontent.com/u/41966024?v=4", "gravatar_id": "", "url": "https://api.github.com/users/nastya236", "html_url": "https://github.com/nastya236", "followers_url": "https://api.github.com/users/nastya236/followers", "following_url": "https://api.github.com/users/nastya236/following{/other_user}", "gists_url": "https://api.github.com/users/nastya236/gists{/gist_id}", "starred_url": "https://api.github.com/users/nastya236/starred{/owner}{/repo}", "subscriptions_url": "https://api.github.com/users/nastya236/subscriptions", "organizations_url": "https://api.github.com/users/nastya236/orgs", "repos_url": "https://api.github.com/users/nastya236/repos", "events_url": "https://api.github.com/users/nastya236/events{/privacy}", "received_events_url": "https://api.github.com/users/nastya236/received_events", "type": "User", "user_view_type": "public", "site_admin": false }, "labels": [], "state": "open", "locked": false, "assignees": [], "milestone": null, "comments": 0, "created_at": "2026-02-02T16:32:39Z", "updated_at": "2026-02-02T18:19:36Z", "closed_at": null, "assignee": null, "author_association": "COLLABORATOR", "issue_field_values": [], "type": null, "active_lock_reason": null, "draft": true, "pull_request": { "url": "https://api.github.com/repos/ml-explore/mlx/pulls/3090", "html_url": "https://github.com/ml-explore/mlx/pull/3090", "diff_url": "https://github.com/ml-explore/mlx/pull/3090.diff", "patch_url": "https://github.com/ml-explore/mlx/pull/3090.patch", "merged_at": null }, "body": "Fuse unary ops into reduction. \r\n\r\nIt will be useful for:\r\n- `mx::max(mx::abs(x))` for `nvfp4` quantization\r\n- `clip_grad_norm`\r\n\r\n```python\r\ndef bench(f, *args, **kwargs):\r\n for i in range(N_warmup):\r\n x = mx.eval(f(*args, **kwargs))\r\n\r\n s = time.perf_counter_ns()\r\n for i in range(N_iter_bench):\r\n x = mx.eval(f(*args, **kwargs))\r\n e = time.perf_counter_ns()\r\n return (e - s) * 1e-9\r\n\r\ndef abs_max(a):\r\n return mx.max(mx.abs(a))\r\n\r\nshape = (4*4096, 11008)\r\nx = mx.random.uniform(shape=shape)\r\ny = mx.random.uniform(shape=shape)\r\nf_com = mx.compile(abs_max)\r\n\r\ntime_fused = bench(f_com, x)\r\ntime_unfused = bench(abs_max, y)\r\n\r\nFused time: 0.001803 s\r\nUnfused time: 0.004267 s\r\n```\r\n\r\nTODO:\r\n- Not sure if we want to: do the same for `row_reduce` and `col_reduce`, seems that we don't do `unary + row/col_reduce` very often \r\n- Metal", "reactions": { "url": "https://api.github.com/repos/ml-explore/mlx/issues/3090/reactions", "total_count": 1, "+1": 0, "-1": 0, "laugh": 0, "hooray": 0, "confused": 0, "heart": 0, "rocket": 1, "eyes": 0 }, "timeline_url": "https://api.github.com/repos/ml-explore/mlx/issues/3090/timeline", "performed_via_github_app": null, "state_reason": null, "score": 1.0 }, { "url": "https://api.github.com/repos/ml-explore/mlx/issues/4058", "repository_url": "https://api.github.com/repos/ml-explore/mlx", "labels_url": "https://api.github.com/repos/ml-explore/mlx/issues/4058/labels{/name}", "comments_url": "https://api.github.com/repos/ml-explore/mlx/issues/4058/comments", "events_url": "https://api.github.com/repos/ml-explore/mlx/issues/4058/events", "html_url": "https://github.com/ml-explore/mlx/pull/4058", "id": 5092426777, "node_id": "PR_kwDOKzRn1878Jd8E", "number": 4058, "title": "Reject a negative max_norm in clip_grad_norm", "user": { "login": "ayaangazali", "id": 231403453, "node_id": "U_kgDODcrvvQ", "avatar_url": "https://avatars.githubusercontent.com/u/231403453?v=4", "gravatar_id": "", "url": "https://api.github.com/users/ayaangazali", "html_url": "https://github.com/ayaangazali", "followers_url": "https://api.github.com/users/ayaangazali/followers", "following_url": "https://api.github.com/users/ayaangazali/following{/other_user}", "gists_url": "https://api.github.com/users/ayaangazali/gists{/gist_id}", "starred_url": "https://api.github.com/users/ayaangazali/starred{/owner}{/repo}", "subscriptions_url": "https://api.github.com/users/ayaangazali/subscriptions", "organizations_url": "https://api.github.com/users/ayaangazali/orgs", "repos_url": "https://api.github.com/users/ayaangazali/repos", "events_url": "https://api.github.com/users/ayaangazali/events{/privacy}", "received_events_url": "https://api.github.com/users/ayaangazali/received_events", "type": "User", "user_view_type": "public", "site_admin": false }, "labels": [], "state": "closed", "locked": false, "assignees": [], "milestone": null, "comments": 0, "created_at": "2026-08-07T17:16:10Z", "updated_at": "2026-08-08T03:13:16Z", "closed_at": "2026-08-08T03:13:16Z", "assignee": null, "author_association": "CONTRIBUTOR", "issue_field_values": [], "type": null, "active_lock_reason": null, "draft": false, "pull_request": { "url": "https://api.github.com/repos/ml-explore/mlx/pulls/4058", "html_url": "https://github.com/ml-explore/mlx/pull/4058", "diff_url": "https://github.com/ml-explore/mlx/pull/4058.diff", "patch_url": "https://github.com/ml-explore/mlx/pull/4058.patch", "merged_at": "2026-08-08T03:13:16Z" }, "body": "## Proposed changes\n\n`clip_grad_norm` accepts a negative `max_norm` and silently reverses every gradient:\n\n```python\n>>> grads = {\"w\": mx.array([0.1, 0.2])}\n>>> opt.clip_grad_norm(grads, -1.0)[0][\"w\"]\narray([-0.447212, -0.894423], dtype=float32)\n```\n\nThe scale factor is `mx.minimum(max_norm / (total_norm + 1e-6), 1.0)`, which is negative whenever `max_norm` is, so the gradients get flipped rather than clipped. Feeding that to an optimizer ascends the loss instead of descending it, with nothing in the output to say so. The docstring promises the opposite: \"It scales down the gradients proportionally if their norm is greater than `max_norm`.\"\n\nThis is reachable without a typo, since `max_norm` is often computed, for example a warmup style `max_norm = base * (1 - step / total)` that goes negative once `step` passes `total`.\n\nEvery other hyperparameter in this module is already range checked and raises, for instance `RMSprop epsilon should be >0, -1 was provided instead` at optimizers.py:614. Added the matching check:\n\n```python\n>>> opt.clip_grad_norm(grads, -1.0)\nValueError: max_norm should be >=0, -1.0 was provided instead\n```\n\n`max_norm=0` is left working and still returns zeroed gradients, since that is a coherent request and not a sign error.\n\n## Checklist\n\nPut an `x` in the boxes that apply.\n\n- [x] I have read the [CONTRIBUTING](https://github.com/ml-explore/mlx/blob/main/CONTRIBUTING.md) document\n- [x] I have run `pre-commit run --all-files` to format my code / installed pre-commit prior to committing changes\n- [x] I have added tests that prove my fix is effective or that my feature works\n- [x] I have updated the necessary documentation (if needed)\n\n`test_optimizers.py` passes (25 tests, CPU only build); the new assertion fails on main.\n\n---\n\nfreshman here, i work alongside Claude Code. found this while feeding out of range hyperparameters to each optimizer and comparing which ones complained. the rest of the module was solid, this was the one that took the bad value and quietly changed the answer.\n", "reactions": { "url": "https://api.github.com/repos/ml-explore/mlx/issues/4058/reactions", "total_count": 0, "+1": 0, "-1": 0, "laugh": 0, "hooray": 0, "confused": 0, "heart": 0, "rocket": 0, "eyes": 0 }, "timeline_url": "https://api.github.com/repos/ml-explore/mlx/issues/4058/timeline", "performed_via_github_app": null, "state_reason": null, "score": 1.0 }, { "url": "https://api.github.com/repos/ml-explore/mlx/issues/1623", "repository_url": "https://api.github.com/repos/ml-explore/mlx", "labels_url": "https://api.github.com/repos/ml-explore/mlx/issues/1623/labels{/name}", "comments_url": "https://api.github.com/repos/ml-explore/mlx/issues/1623/comments", "events_url": "https://api.github.com/repos/ml-explore/mlx/issues/1623/events", "html_url": "https://github.com/ml-explore/mlx/issues/1623", "id": 2689278465, "node_id": "I_kwDOKzRn186gSyIB", "number": 1623, "title": "[COMMENT] AdamW including bias + Clip grad norm closer to pytorch performances", "user": { "login": "thegodone", "id": 1186658, "node_id": "MDQ6VXNlcjExODY2NTg=", "avatar_url": "https://avatars.githubusercontent.com/u/1186658?v=4", "gravatar_id": "", "url": "https://api.github.com/users/thegodone", "html_url": "https://github.com/thegodone", "followers_url": "https://api.github.com/users/thegodone/followers", "following_url": "https://api.github.com/users/thegodone/following{/other_user}", "gists_url": "https://api.github.com/users/thegodone/gists{/gist_id}", "starred_url": "https://api.github.com/users/thegodone/starred{/owner}{/repo}", "subscriptions_url": "https://api.github.com/users/thegodone/subscriptions", "organizations_url": "https://api.github.com/users/thegodone/orgs", "repos_url": "https://api.github.com/users/thegodone/repos", "events_url": "https://api.github.com/users/thegodone/events{/privacy}", "received_events_url": "https://api.github.com/users/thegodone/received_events", "type": "User", "user_view_type": "public", "site_admin": false }, "labels": [], "state": "closed", "locked": true, "assignees": [], "milestone": null, "comments": 0, "created_at": "2024-11-25T05:49:36Z", "updated_at": "2026-04-14T10:19:06Z", "closed_at": "2026-04-14T10:19:06Z", "assignee": null, "author_association": "NONE", "issue_field_values": [], "type": null, "active_lock_reason": null, "sub_issues_summary": { "total": 0, "completed": 0, "percent_completed": 0 }, "issue_dependencies_summary": { "blocked_by": 0, "total_blocked_by": 0, "blocking": 0, "total_blocking": 0 }, "body": "**Comment:**\r\n\r\nUsing the AdamW with bias correction (not the one in mlx) combine with clip_grad_norm allows to be similar to pytorch performances. \r\n\r\nOne possible explanation is that pytorch also has a trick to avoid gradient explosion. \r\n\r\nAnother one explanation, more logical, is that the initialisation has a big impact on gradient stability. \r\n\r\nThe fact that the LR changes at each iteration instead of each epoch can be the last one.\r\n\r\nSuggestion feature : Having LR scheduler with a option to do it per batch or per epoch would be a great feature to evaluate. \r\n\r\nHere a view of 10 time runs using the same initialization (loading identical saved weights for each run) & same data split (mv = validation, mt = test)\r\n\r\n![image](https://github.com/user-attachments/assets/9741540e-6aa8-4e4a-b5d2-107b99583920)\r\n.\r\nwe can see the real deviation of RMSE over the runs.\r\n\r\n**Desktop (please complete the following information):**\r\n - OS Version: [e.g. MacOS 15.1.1]\r\n - Version [e.g. 0.21.0]\r\n", "reactions": { "url": "https://api.github.com/repos/ml-explore/mlx/issues/1623/reactions", "total_count": 0, "+1": 0, "-1": 0, "laugh": 0, "hooray": 0, "confused": 0, "heart": 0, "rocket": 0, "eyes": 0 }, "timeline_url": "https://api.github.com/repos/ml-explore/mlx/issues/1623/timeline", "performed_via_github_app": null, "state_reason": "completed", "pinned_comment": null, "score": 1.0 }, { "url": "https://api.github.com/repos/ml-explore/mlx/issues/2837", "repository_url": "https://api.github.com/repos/ml-explore/mlx", "labels_url": "https://api.github.com/repos/ml-explore/mlx/issues/2837/labels{/name}", "comments_url": "https://api.github.com/repos/ml-explore/mlx/issues/2837/comments", "events_url": "https://api.github.com/repos/ml-explore/mlx/issues/2837/events", "html_url": "https://github.com/ml-explore/mlx/issues/2837", "id": 3665597255, "node_id": "I_kwDOKzRn187afJNH", "number": 2837, "title": "Feature Request: Add support for parameter groups in optimizers", "user": { "login": "yuchaoran2011", "id": 1168769, "node_id": "MDQ6VXNlcjExNjg3Njk=", "avatar_url": "https://avatars.githubusercontent.com/u/1168769?v=4", "gravatar_id": "", "url": "https://api.github.com/users/yuchaoran2011", "html_url": "https://github.com/yuchaoran2011", "followers_url": "https://api.github.com/users/yuchaoran2011/followers", "following_url": "https://api.github.com/users/yuchaoran2011/following{/other_user}", "gists_url": "https://api.github.com/users/yuchaoran2011/gists{/gist_id}", "starred_url": "https://api.github.com/users/yuchaoran2011/starred{/owner}{/repo}", "subscriptions_url": "https://api.github.com/users/yuchaoran2011/subscriptions", "organizations_url": "https://api.github.com/users/yuchaoran2011/orgs", "repos_url": "https://api.github.com/users/yuchaoran2011/repos", "events_url": "https://api.github.com/users/yuchaoran2011/events{/privacy}", "received_events_url": "https://api.github.com/users/yuchaoran2011/received_events", "type": "User", "user_view_type": "public", "site_admin": false }, "labels": [], "state": "closed", "locked": false, "assignees": [], "milestone": null, "comments": 2, "created_at": "2025-11-26T03:46:27Z", "updated_at": "2025-11-27T08:37:57Z", "closed_at": "2025-11-27T08:37:57Z", "assignee": null, "author_association": "CONTRIBUTOR", "issue_field_values": [], "type": null, "active_lock_reason": null, "sub_issues_summary": { "total": 0, "completed": 0, "percent_completed": 0 }, "issue_dependencies_summary": { "blocked_by": 0, "total_blocked_by": 0, "blocking": 0, "total_blocking": 0 }, "body": "Currently, MLX optimizers only accept scalar hyperparameters, requiring all parameters to use the same learning rate, weight decay, etc. This differs from PyTorch's optimizer API, which supports parameter groups with different configurations.\n```\noptimizer = optim.AdamW(learning_rate=1e-4, weight_decay=0.01)\n```\nThis means all parameters must use identical optimizer settings, which is problematic for common training practices like Applying weight decay only to weight matrices but not biases. \n\nI had to create multiple optimizers and manually split the gradients into decay groups and no-decay groups every training step. This approach introduces code complexity and is less efficient:\n```\noptimizer_decay = optim.AdamW(learning_rate=1e-4, weight_decay=0.01)\noptimizer_skip_decay = optim.AdamW(learning_rate=1e-3, weight_decay=0.0)\ndef step(x, y):\n loss, grads = ...\n # Split gradients into two groups\n grads_to_decay, grads_to_skip_decay = {}, {}\n split_grads(grads, grads_to_decay, grads_to_skip_decay) # Complex recursive logic\n \n # Clip and update separately\n clipped_grads_to_decay, _ = clip_grad_norm(grads_to_decay, max_norm=1.0)\n clipped_grads_to_skip_decay, _ = clip_grad_norm(grads_to_skip_decay, max_norm=1.0)\n optimizer_decay.update(model, clipped_grads_to_decay)\n optimizer_skip_decay.update(model, clipped_grads_to_skip_decay)\n```\nIf we have support for parameter group like in PyTorch, the above code can be simplified to:\n```\n# Single optimizer with parameter groups\noptimizer = optim.AdamW([\n {'params': decay_params, 'learning_rate': 1e-4, 'weight_decay': 0.01},\n {'params': no_decay_params, 'learning_rate': 1e-3, 'weight_decay': 0.0}\n])\n\ndef step(x, y):\n loss, grads = ...\n clipped_grads, _ = clip_grad_norm(grads, max_norm=1.0)\n optimizer.update(model, clipped_grads) \n```\nPlease consider this improvement proposal. If there's any simpler way of doing it in MLX that I missed, feel free to suggest.", "reactions": { "url": "https://api.github.com/repos/ml-explore/mlx/issues/2837/reactions", "total_count": 0, "+1": 0, "-1": 0, "laugh": 0, "hooray": 0, "confused": 0, "heart": 0, "rocket": 0, "eyes": 0 }, "timeline_url": "https://api.github.com/repos/ml-explore/mlx/issues/2837/timeline", "performed_via_github_app": null, "state_reason": "completed", "pinned_comment": null, "score": 1.0 }, { "url": "https://api.github.com/repos/ml-explore/mlx/issues/1622", "repository_url": "https://api.github.com/repos/ml-explore/mlx", "labels_url": "https://api.github.com/repos/ml-explore/mlx/issues/1622/labels{/name}", "comments_url": "https://api.github.com/repos/ml-explore/mlx/issues/1622/comments", "events_url": "https://api.github.com/repos/ml-explore/mlx/issues/1622/events", "html_url": "https://github.com/ml-explore/mlx/issues/1622", "id": 2688278143, "node_id": "I_kwDOKzRn186gO95_", "number": 1622, "title": "[QUESTION] how to use clip_grad_norm ?", "user": { "login": "thegodone", "id": 1186658, "node_id": "MDQ6VXNlcjExODY2NTg=", "avatar_url": "https://avatars.githubusercontent.com/u/1186658?v=4", "gravatar_id": "", "url": "https://api.github.com/users/thegodone", "html_url": "https://github.com/thegodone", "followers_url": "https://api.github.com/users/thegodone/followers", "following_url": "https://api.github.com/users/thegodone/following{/other_user}", "gists_url": "https://api.github.com/users/thegodone/gists{/gist_id}", "starred_url": "https://api.github.com/users/thegodone/starred{/owner}{/repo}", "subscriptions_url": "https://api.github.com/users/thegodone/subscriptions", "organizations_url": "https://api.github.com/users/thegodone/orgs", "repos_url": "https://api.github.com/users/thegodone/repos", "events_url": "https://api.github.com/users/thegodone/events{/privacy}", "received_events_url": "https://api.github.com/users/thegodone/received_events", "type": "User", "user_view_type": "public", "site_admin": false }, "labels": [], "state": "closed", "locked": false, "assignees": [], "milestone": null, "comments": 1, "created_at": "2024-11-24T19:34:44Z", "updated_at": "2024-11-25T00:41:59Z", "closed_at": "2024-11-25T00:41:59Z", "assignee": null, "author_association": "NONE", "issue_field_values": [], "type": null, "active_lock_reason": null, "sub_issues_summary": { "total": 0, "completed": 0, "percent_completed": 0 }, "issue_dependencies_summary": { "blocked_by": 0, "total_blocked_by": 0, "blocking": 0, "total_blocking": 0 }, "body": "**Describe the question**\r\n\r\nHow to use the clp_grap_norm in training model ?\r\n\r\n**To Reproduce**\r\n\r\nInclude code snippet\r\n````\r\nclipped_grads, total_norm = optim.clip_grad_norm(grads, max_norm=5.0)\r\n````\r\n\r\njust like this ?\r\n\r\n````\r\ndef step(model, graph, labels):\r\n loss_and_grad_fn = nn.value_and_grad(model, forward_fn)\r\n (loss, y_hat), grads = loss_and_grad_fn(\r\n model=model,\r\n graph=graph,\r\n labels=labels,\r\n )\r\n clipped_grads, total_norm = optim.clip_grad_norm(grads, max_norm=5.0) \r\n optimizer.update(model, clipped_grads)\r\n return loss\r\n\r\n````\r\n\r\n**Desktop (please complete the following information):**\r\n - OS Version: [e.g. MacOS 15.1.1]\r\n - Version [e.g. 0.21]\r\n", "reactions": { "url": "https://api.github.com/repos/ml-explore/mlx/issues/1622/reactions", "total_count": 0, "+1": 0, "-1": 0, "laugh": 0, "hooray": 0, "confused": 0, "heart": 0, "rocket": 0, "eyes": 0 }, "timeline_url": "https://api.github.com/repos/ml-explore/mlx/issues/1622/timeline", "performed_via_github_app": null, "state_reason": "completed", "pinned_comment": null, "score": 1.0 }, { "url": "https://api.github.com/repos/ml-explore/mlx/issues/1043", "repository_url": "https://api.github.com/repos/ml-explore/mlx", "labels_url": "https://api.github.com/repos/ml-explore/mlx/issues/1043/labels{/name}", "comments_url": "https://api.github.com/repos/ml-explore/mlx/issues/1043/comments", "events_url": "https://api.github.com/repos/ml-explore/mlx/issues/1043/events", "html_url": "https://github.com/ml-explore/mlx/pull/1043", "id": 2266965189, "node_id": "PR_kwDOKzRn185t58Nd", "number": 1043, "title": "feat: implement `clip_grad_norm`", "user": { "login": "NripeshN", "id": 86844847, "node_id": "MDQ6VXNlcjg2ODQ0ODQ3", "avatar_url": "https://avatars.githubusercontent.com/u/86844847?v=4", "gravatar_id": "", "url": "https://api.github.com/users/NripeshN", "html_url": "https://github.com/NripeshN", "followers_url": "https://api.github.com/users/NripeshN/followers", "following_url": "https://api.github.com/users/NripeshN/following{/other_user}", "gists_url": "https://api.github.com/users/NripeshN/gists{/gist_id}", "starred_url": "https://api.github.com/users/NripeshN/starred{/owner}{/repo}", "subscriptions_url": "https://api.github.com/users/NripeshN/subscriptions", "organizations_url": "https://api.github.com/users/NripeshN/orgs", "repos_url": "https://api.github.com/users/NripeshN/repos", "events_url": "https://api.github.com/users/NripeshN/events{/privacy}", "received_events_url": "https://api.github.com/users/NripeshN/received_events", "type": "User", "user_view_type": "public", "site_admin": false }, "labels": [], "state": "closed", "locked": false, "assignees": [], "milestone": null, "comments": 1, "created_at": "2024-04-27T11:02:23Z", "updated_at": "2024-05-03T17:02:08Z", "closed_at": "2024-05-03T16:07:02Z", "assignee": null, "author_association": "CONTRIBUTOR", "issue_field_values": [], "type": null, "active_lock_reason": null, "draft": false, "pull_request": { "url": "https://api.github.com/repos/ml-explore/mlx/pulls/1043", "html_url": "https://github.com/ml-explore/mlx/pull/1043", "diff_url": "https://github.com/ml-explore/mlx/pull/1043.diff", "patch_url": "https://github.com/ml-explore/mlx/pull/1043.patch", "merged_at": "2024-05-03T16:07:02Z" }, "body": "## Proposed changes\r\n\r\nCloses #1040\r\nCloses https://github.com/ml-explore/mlx/issues/1045\r\n\r\n## Checklist\r\n\r\nPut an `x` in the boxes that apply.\r\n\r\n- [x] I have read the [CONTRIBUTING](https://github.com/ml-explore/mlx/blob/main/CONTRIBUTING.md) document\r\n- [x] I have run `pre-commit run --all-files` to format my code / installed pre-commit prior to committing changes\r\n- [x] I have added tests that prove my fix is effective or that my feature works\r\n- [x] I have updated the necessary documentation (if needed)\r\n", "reactions": { "url": "https://api.github.com/repos/ml-explore/mlx/issues/1043/reactions", "total_count": 2, "+1": 2, "-1": 0, "laugh": 0, "hooray": 0, "confused": 0, "heart": 0, "rocket": 0, "eyes": 0 }, "timeline_url": "https://api.github.com/repos/ml-explore/mlx/issues/1043/timeline", "performed_via_github_app": null, "state_reason": null, "score": 1.0 }, { "url": "https://api.github.com/repos/ml-explore/mlx/issues/1040", "repository_url": "https://api.github.com/repos/ml-explore/mlx", "labels_url": "https://api.github.com/repos/ml-explore/mlx/issues/1040/labels{/name}", "comments_url": "https://api.github.com/repos/ml-explore/mlx/issues/1040/comments", "events_url": "https://api.github.com/repos/ml-explore/mlx/issues/1040/events", "html_url": "https://github.com/ml-explore/mlx/issues/1040", "id": 2266514803, "node_id": "I_kwDOKzRn186HGEVz", "number": 1040, "title": "[Feature] Add a `clip_grad_norm` to `mlx.optimizers`", "user": { "login": "awni", "id": 1542805, "node_id": "MDQ6VXNlcjE1NDI4MDU=", "avatar_url": "https://avatars.githubusercontent.com/u/1542805?v=4", "gravatar_id": "", "url": "https://api.github.com/users/awni", "html_url": "https://github.com/awni", "followers_url": "https://api.github.com/users/awni/followers", "following_url": "https://api.github.com/users/awni/following{/other_user}", "gists_url": "https://api.github.com/users/awni/gists{/gist_id}", "starred_url": "https://api.github.com/users/awni/starred{/owner}{/repo}", "subscriptions_url": "https://api.github.com/users/awni/subscriptions", "organizations_url": "https://api.github.com/users/awni/orgs", "repos_url": "https://api.github.com/users/awni/repos", "events_url": "https://api.github.com/users/awni/events{/privacy}", "received_events_url": "https://api.github.com/users/awni/received_events", "type": "User", "user_view_type": "public", "site_admin": false }, "labels": [ { "id": 6264100360, "node_id": "LA_kwDOKzRn188AAAABdV6WCA", "url": "https://api.github.com/repos/ml-explore/mlx/labels/enhancement", "name": "enhancement", "color": "a2eeef", "default": true, "description": "" }, { "id": 6264100366, "node_id": "LA_kwDOKzRn188AAAABdV6WDg", "url": "https://api.github.com/repos/ml-explore/mlx/labels/good%20first%20issue", "name": "good first issue", "color": "7057ff", "default": true, "description": "" } ], "state": "closed", "locked": false, "assignees": [], "milestone": null, "comments": 0, "created_at": "2024-04-26T21:12:58Z", "updated_at": "2024-05-03T16:07:03Z", "closed_at": "2024-05-03T16:07:03Z", "assignee": null, "author_association": "MEMBER", "issue_field_values": [], "type": null, "active_lock_reason": null, "sub_issues_summary": { "total": 0, "completed": 0, "percent_completed": 0 }, "issue_dependencies_summary": { "blocked_by": 0, "total_blocked_by": 0, "blocking": 0, "total_blocking": 0 }, "body": "It's relatively simple to implement, but maybe worth adding since it's also quite common. Here's a possible implementation:\r\n\r\n```python\r\nimport mlx.core as mx\r\nfrom mlx.utils import tree_map\r\n\r\ndef clip_grad_norm(grads, max_norm):\r\n norm = mx.array(0.0)\r\n def accumulate(g):\r\n nonlocal norm\r\n norm += g.square().sum()\r\n\r\n tree_map(accumulate, grads)\r\n norm = mx.sqrt(norm)\r\n\r\n clip = lambda g: mx.where(norm < max_norm, g, g * max_norm / (norm + 1e-6))\r\n grads = tree_map(clip, grads)\r\n return grads, norm\r\n```\r\n\r\nWe could also add a `tree_reduce` which would simplify it a bit.", "reactions": { "url": "https://api.github.com/repos/ml-explore/mlx/issues/1040/reactions", "total_count": 1, "+1": 1, "-1": 0, "laugh": 0, "hooray": 0, "confused": 0, "heart": 0, "rocket": 0, "eyes": 0 }, "timeline_url": "https://api.github.com/repos/ml-explore/mlx/issues/1040/timeline", "performed_via_github_app": null, "state_reason": "completed", "pinned_comment": null, "score": 1.0 } ], "search_type": "lexical" } }, { "query": "repo:ml-explore/mlx \"gradient\" \"clipping\" \"overflow\"", "url": "https://api.github.com/search/issues?q=repo%3Aml-explore%2Fmlx+%22gradient%22+%22clipping%22+%22overflow%22&per_page=100", "result": { "total_count": 0, "incomplete_results": false, "items": [], "search_type": "lexical" } }, { "query": "repo:ml-explore/mlx \"norm\" \"overflow\"", "url": "https://api.github.com/search/issues?q=repo%3Aml-explore%2Fmlx+%22norm%22+%22overflow%22&per_page=100&sort=updated", "result": { "total_count": 2, "incomplete_results": false, "items": [ { "url": "https://api.github.com/repos/ml-explore/mlx/issues/4230", "repository_url": "https://api.github.com/repos/ml-explore/mlx", "labels_url": "https://api.github.com/repos/ml-explore/mlx/issues/4230/labels{/name}", "comments_url": "https://api.github.com/repos/ml-explore/mlx/issues/4230/comments", "events_url": "https://api.github.com/repos/ml-explore/mlx/issues/4230/events", "html_url": "https://github.com/ml-explore/mlx/pull/4230", "id": 5139917446, "node_id": "PR_kwDOKzRn187-gNam", "number": 4230, "title": "Stabilize reduced-precision InstanceNorm", "user": { "login": "ternaus", "id": 5481618, "node_id": "MDQ6VXNlcjU0ODE2MTg=", "avatar_url": "https://avatars.githubusercontent.com/u/5481618?v=4", "gravatar_id": "", "url": "https://api.github.com/users/ternaus", "html_url": "https://github.com/ternaus", "followers_url": "https://api.github.com/users/ternaus/followers", "following_url": "https://api.github.com/users/ternaus/following{/other_user}", "gists_url": "https://api.github.com/users/ternaus/gists{/gist_id}", "starred_url": "https://api.github.com/users/ternaus/starred{/owner}{/repo}", "subscriptions_url": "https://api.github.com/users/ternaus/subscriptions", "organizations_url": "https://api.github.com/users/ternaus/orgs", "repos_url": "https://api.github.com/users/ternaus/repos", "events_url": "https://api.github.com/users/ternaus/events{/privacy}", "received_events_url": "https://api.github.com/users/ternaus/received_events", "type": "User", "user_view_type": "public", "site_admin": false }, "labels": [ { "id": 6381853406, "node_id": "LA_kwDOKzRn188AAAABfGNa3g", "url": "https://api.github.com/repos/ml-explore/mlx/labels/low%20priority", "name": "low priority", "color": "c2e0c6", "default": false, "description": "" } ], "state": "closed", "locked": false, "assignees": [], "milestone": null, "comments": 0, "created_at": "2026-08-13T09:34:29Z", "updated_at": "2026-08-22T01:49:57Z", "closed_at": "2026-08-22T01:49:57Z", "assignee": null, "author_association": "CONTRIBUTOR", "issue_field_values": [], "type": null, "active_lock_reason": null, "draft": false, "pull_request": { "url": "https://api.github.com/repos/ml-explore/mlx/pulls/4230", "html_url": "https://github.com/ml-explore/mlx/pull/4230", "diff_url": "https://github.com/ml-explore/mlx/pull/4230.diff", "patch_url": "https://github.com/ml-explore/mlx/pull/4230.patch", "merged_at": "2026-08-22T01:49:57Z" }, "body": "## Proposed changes\n\nFixes #4228.\n\nThis PR stabilizes reduced-precision `InstanceNorm` without changing its public `(N, ..., C)` layout or output dtype:\n\n- reshape each input to `(N, C, spatial_voxels)` and use the stable fused `mx.fast.layer_norm` reduction;\n- restore the original channels-last shape before the optional affine transform;\n- add an FP16 regression case whose direct variance overflows even though every input value is finite.\n\n## Why\n\nThe current FP16 path computes `mx.var` in FP16. A channel alternating between `-512` and `512` therefore gets `var=inf`, and `InstanceNorm` collapses the channel to zero instead of returning values near `-1` and `1`.\n\n`mx.fast.layer_norm` already performs the same reduction needed by `InstanceNorm` after the spatial dimensions are flattened. It avoids a full FP32 input cast and uses the optimized normalization kernel.\n\n## Performance\n\nApple M4, MLX 0.32.0 core, FP16 input, five warmups and 20 measured calls:\n\n| NDHWC shape | Existing p50 | FP32-cast p50 | This PR p50 |\n| --- | ---: | ---: | ---: |\n| `1 x 32 x 64 x 64 x 30` | `0.785 ms` | `0.789 ms` | `0.196 ms` |\n| `1 x 128 x 96 x 96 x 30` | `2.357 ms` | `5.145 ms` | `1.098 ms` |\n\nThe larger shape is the first-stage feature map from the 3D segmentation route that exposed the overflow.\n\n## Validation\n\n- `PYTHONPATH=python:python/tests python -m unittest test_nn.TestLayers`\n- `PYTHONPATH=python:python/tests python -m unittest test_nn`\n- `pre-commit run --all-files`\n\n## Checklist\n\n- [x] I have read the [CONTRIBUTING](https://github.com/ml-explore/mlx/blob/main/CONTRIBUTING.md) document\n- [x] I have run `pre-commit run --all-files` to format my code / installed pre-commit prior to committing changes\n- [x] I have added tests that prove my fix is effective or that my feature works\n- [x] I have updated the necessary documentation (if needed; no public API or documented behavior changed)\n", "reactions": { "url": "https://api.github.com/repos/ml-explore/mlx/issues/4230/reactions", "total_count": 1, "+1": 1, "-1": 0, "laugh": 0, "hooray": 0, "confused": 0, "heart": 0, "rocket": 0, "eyes": 0 }, "timeline_url": "https://api.github.com/repos/ml-explore/mlx/issues/4230/timeline", "performed_via_github_app": null, "state_reason": null, "score": 1.0 }, { "url": "https://api.github.com/repos/ml-explore/mlx/issues/1049", "repository_url": "https://api.github.com/repos/ml-explore/mlx", "labels_url": "https://api.github.com/repos/ml-explore/mlx/issues/1049/labels{/name}", "comments_url": "https://api.github.com/repos/ml-explore/mlx/issues/1049/comments", "events_url": "https://api.github.com/repos/ml-explore/mlx/issues/1049/events", "html_url": "https://github.com/ml-explore/mlx/issues/1049", "id": 2267671610, "node_id": "I_kwDOKzRn186HKew6", "number": 1049, "title": "[FEATURE] in keras LayerNorm by default is apply to last dimension only", "user": { "login": "thegodone", "id": 1186658, "node_id": "MDQ6VXNlcjExODY2NTg=", "avatar_url": "https://avatars.githubusercontent.com/u/1186658?v=4", "gravatar_id": "", "url": "https://api.github.com/users/thegodone", "html_url": "https://github.com/thegodone", "followers_url": "https://api.github.com/users/thegodone/followers", "following_url": "https://api.github.com/users/thegodone/following{/other_user}", "gists_url": "https://api.github.com/users/thegodone/gists{/gist_id}", "starred_url": "https://api.github.com/users/thegodone/starred{/owner}{/repo}", "subscriptions_url": "https://api.github.com/users/thegodone/subscriptions", "organizations_url": "https://api.github.com/users/thegodone/orgs", "repos_url": "https://api.github.com/users/thegodone/repos", "events_url": "https://api.github.com/users/thegodone/events{/privacy}", "received_events_url": "https://api.github.com/users/thegodone/received_events", "type": "User", "user_view_type": "public", "site_admin": false }, "labels": [], "state": "closed", "locked": false, "assignees": [], "milestone": null, "comments": 9, "created_at": "2024-04-28T16:35:36Z", "updated_at": "2024-04-29T14:39:17Z", "closed_at": "2024-04-28T17:23:28Z", "assignee": null, "author_association": "NONE", "issue_field_values": [], "type": null, "active_lock_reason": null, "sub_issues_summary": { "total": 0, "completed": 0, "percent_completed": 0 }, "issue_dependencies_summary": { "blocked_by": 0, "total_blocked_by": 0, "blocking": 0, "total_blocking": 0 }, "body": "**Describe the bug**\r\nLook like mlx.core.LayerNorm applies on all dimensions. Can we add an axis parameter to set it like in keras ?\r\n\r\n**Expected behavior**\r\nI am trying to clone a keras code and the layernorm behaviour is not identical to keras. Can we add a feature to apply only on the axis we want to.\r\n\r\n\r\n**Desktop (please complete the following information):**\r\n - OS Version: MacOS 14.4\r\n - Version 0.7.0\r\n\r\n**Additional context**\r\nAdd any other context about the problem here.\r\n", "reactions": { "url": "https://api.github.com/repos/ml-explore/mlx/issues/1049/reactions", "total_count": 0, "+1": 0, "-1": 0, "laugh": 0, "hooray": 0, "confused": 0, "heart": 0, "rocket": 0, "eyes": 0 }, "timeline_url": "https://api.github.com/repos/ml-explore/mlx/issues/1049/timeline", "performed_via_github_app": null, "state_reason": "completed", "pinned_comment": null, "score": 1.0 } ], "search_type": "lexical" } }, { "query": "repo:ml-explore/mlx \"linalg.norm\" \"underflow\"", "url": "https://api.github.com/search/issues?q=repo%3Aml-explore%2Fmlx+%22linalg.norm%22+%22underflow%22&per_page=100", "result": { "total_count": 0, "incomplete_results": false, "items": [], "search_type": "lexical" } } ]