45af0df90b
Fix repr of conv layers ( #3275 )
2026-03-18 22:47:38 -07:00
f226eeec9e
Fix nn.GRU skipping bhn bias when hidden is None ( #3252 )
...
Co-authored-by: mm65x <mm65x@users.noreply.github.com >
2026-03-16 13:28:14 -07:00
e1e1399e1b
Hybrid sharding ( #3194 )
2026-03-10 11:47:25 +01:00
d2702a4fc1
Fix non-strict module update with extra weights ( #3214 )
...
Co-authored-by: Michelle DiMarco <m_dimarco@apple.com >
2026-03-09 22:03:50 -07:00
d06c3c8936
Improve mlx.distributed_config ( #3199 )
2026-03-09 13:17:51 -07:00
72e04f7fb7
[CUDA] Fsdp (easy) ( #3130 )
2026-03-01 23:29:09 +01:00
e226af720e
Propagate quantization mode in quantized layers ( #3133 )
2026-02-15 18:33:13 -08:00
5e018de4e5
Quantize module to QQLinear ( #3106 )
2026-02-09 14:35:17 -08:00
ceea571490
JACCL update ( #3094 )
2026-02-05 15:16:07 -08:00
590b4f1c16
Fix ALiBi slopes for non-power-of-2 num_heads ( #3071 )
2026-01-29 07:23:11 -08:00
c86a9bced1
[Docs] Simple example of using MLX distributed ( #2973 )
...
Co-authored-by: Awni Hannun <awni@apple.com >
2026-01-28 17:14:56 -08:00
46d0fdc5ec
Type Enhancement for Func Transforms and Bug Fix ( #3003 )
...
Co-authored-by: Awni Hannun <awni@apple.com >
2026-01-20 07:19:57 +09:00
ca14d3d835
Fix sharding of quantized models with non-power-of-2 bits ( #3006 )
2026-01-18 07:21:56 -08:00
d2bef3c6bb
fix distributed all_to_sharded bias shard axis from -2 to -1 ( #2987 )
2026-01-17 06:51:42 -08:00
3fe7794f22
Reverts changing the MLX_IBV_DEVICES to MLX_JACCL_DEVICES ( #2999 )
2026-01-14 15:44:17 -08:00
a8197795f5
replace MLX_IBV_COORDINATOR with MLX_JACCL_COORDINATOR ( #2986 )
2026-01-13 11:26:25 -08:00
503731727d
QQ linear ( #2931 )
2026-01-05 11:20:54 -08:00
1df6c2a009
Fix doc issues in mlx.nn.init.he_normal and mlx.nn.hard_tanh ( #2968 )
...
Co-authored-by: Awni Hannun <awni@apple.com >
2026-01-05 07:23:41 -08:00
1eef1d155c
Metal/CPU nvfp4 and mxfp8 ( #2946 )
2025-12-22 20:45:19 -08:00
9cfda1a86e
Fixes in mlx.distributed_config ( #2947 )
2025-12-22 17:38:52 -08:00
d9f4d8d508
Fix pid in local launch ( #2936 )
2025-12-19 13:09:15 -08:00
9194ec20a8
Thunderbolt RDMA communications backend ( #2808 )
2025-12-17 11:27:54 -08:00
9abb0b8123
Added support for pytree types that inherit from tuple and typing.namedtuple ( #2845 )
2025-12-04 11:06:45 -08:00
b054838780
Added clarification to apply_fn parameter of apply_to_modules ( #2831 )
...
Co-authored-by: Awni Hannun <awni@apple.com >
2025-11-26 15:40:56 -08:00
eba6a9d163
Compatibility with pip-installed openmpi ( #2741 )
2025-11-07 16:58:31 -08:00
26ceb507eb
only build for macos 14 and up ( #2731 )
...
* only build for macos 14 and up
* bump metal cpp
2025-11-04 09:44:15 -08:00
b901a9f311
Fix the order of hosts in the ring ( #2718 )
2025-10-30 15:02:39 -07:00
793a31eeb6
Fix missing domain_uuid_key in thunderbolt ring setup ( #2682 )
2025-10-30 13:17:20 -07:00
e88f2d4a8e
fix cross entropy axis param ( #2641 )
...
* fix cross entropy axis param
* faster grad clipping
2025-10-01 16:49:55 -07:00
db5443e831
Adding Relu2 ( #2582 )
...
* in. com.
* upd. ackn.
* update __init__
* nits
* nits + format
* used mx.maximum(x, 0) instead of calling the function and moves relu6 under relu2 to make it nicer
* same with _make_activation_module
* Update python/mlx/nn/layers/activations.py
upd
Co-authored-by: Awni Hannun <awni.hannun@gmail.com >
* update funct.rst
* upd. layers.rst
---------
Co-authored-by: Awni Hannun <awni.hannun@gmail.com >
2025-09-10 07:24:30 -07:00
8f163a367d
typing: add type hints to mlx.core.array, linalg, distributed, and random ( #2565 )
...
* Add type annotations to mlx methods
* Missing list_or_scalar
2025-09-04 09:08:11 -07:00
89a3df9014
Fixed several type annotations in the MLX stubs which degraded to Unknown/Any ( #2560 )
...
* Added scalar to stubs to fix Unkown Type Hint
### Proposed changes
Issue #2478 reports that several type annotations in the MLX stubs degrade to Unknown/Any in editors like VS Code with Pylance, due to missing imports (Union, Optional, Tuple) and an undefined scalar type alias.
This PR updates the stub generation patterns to:
• Add missing typing imports in mlx.core.__prefix__ so that Union, Optional, Tuple, etc. are always available.
• Define and export scalar: TypeAlias = Union[int, float, bool] in mlx.core.__suffix__ so that functions typed with Union[scalar, array] resolve correctly instead of falling back to Any.
• Update submodule stub prefixes (distributed, fast, linalg, metal, random) to import scalar alongside array, Device, and Stream, ensuring type checkers resolve the union consistently across modules.
With these changes, functions like mlx.add now display rich type signatures such as:
```
def add(
a: scalar | array,
b: scalar | array,
stream: Stream | Device | None = None
) -> array
```
instead of degrading to Any.
### Checklist
• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works (n/a — stub generation only)
• I have updated the necessary documentation (if needed)
* add bool to patterns
---------
Co-authored-by: Awni Hannun <awni@apple.com >
2025-09-03 12:52:08 -07:00
c5460762e7
Fix AdamW weight_decay default value in docstring ( #2557 )
2025-08-31 21:29:30 -07:00
111f1e71af
Faster contiguous gather for indices in the first axis ( #2552 )
...
* faster contiguous gather for indices in the first axis
* work per thread > 1
* angelos suggestion for scales / biases
2025-08-28 21:26:30 -07:00
70560b6bd5
Add mode parameter for quantization ( #2499 )
...
* add mode parameter for quantization
* mxfp4 quantize/dequantize + start of optional biases
* mxfp4 works
* speedup
* cpu mxfp4
* fix
* fix test tol
* fix
* refactor
* add quant mode enum
2025-08-28 06:45:26 -07:00
3dcb286baf
Remove stream from average grads so it uses default ( #2532 )
...
* Remove stream from average grads so it uses default
* comment
2025-08-25 15:56:29 -07:00
068a4612e9
nccl default for backend=any ( #2528 )
...
* nccl default for backend=any
* check num gpus + ensure row contiguous for all reduce
* comment
2025-08-22 12:24:27 -07:00
f93f87c802
nccl dep + default for cuda ( #2526 )
2025-08-21 17:57:49 -07:00
9392fc3f88
NCCL backend ( #2476 )
2025-08-21 11:56:15 -07:00
728d4db582
Support destination arg in tree flatten/unflatten ( #2450 )
2025-08-06 15:34:59 -07:00
b405591249
fix circular reference ( #2443 )
2025-07-30 09:37:44 -07:00
7d9d6ef456
docs: fix adam and adamw eps placement ( #2416 )
...
Co-authored-by: Mikhail Gorbunov <m_gorbunov@apple.com >
2025-07-24 16:40:45 -07:00
deee214a95
Adding support for the Muon Optimizer ( #1914 )
...
* initial commit with workong optmimizer
* update ACKNOWLEDGMENTS.md
* nits and adding it to test
* nits
* G.astype(mx.bfloat16) to G.astype(G.dtype)
* G.ndim >= 2 to assert G.ndim == 2
* remove coments
* replace with mx.addmm
* remove comments
* format
* nits
* match muon
* fix addmm
---------
Co-authored-by: Awni Hannun <awni@apple.com >
2025-07-18 12:25:28 -07:00
0eb035b4b1
Fix type promotion in Adam with bias correction ( #2350 )
2025-07-10 11:14:42 -07:00
19facd4b20
Build with all cpu cores by default ( #2336 )
2025-07-07 06:06:45 -07:00
cfb6a244ea
allow parameters to be deleted ( #2325 )
2025-07-01 21:27:23 -07:00
33bf1a244b
Fix module update in strict mode ( #2321 )
...
* fix module update in strict mode
* allow GELU to be pickled
2025-06-29 11:12:29 -07:00
5adf185f86
Fix update_modules() when providing a subset ( #2308 )
2025-06-20 17:19:46 -07:00
004c1d8ef2
Report number of missing parameters ( #2264 )
...
* chore: inform
* chore: format
---------
Co-authored-by: FL33TW00D <FL33TW00D@users.noreply.github.com >
2025-06-10 06:37:50 -07:00
c763fe1be0
default strict mode for module update and update_modules ( #2239 )
2025-06-05 15:27:02 -07:00
Angelos Katharopoulos