a8776b7bbd
Fix: Correct cross-attention query routing in Post-LN TransformerDecoderLayer ( #3382 )
2026-04-07 09:16:12 -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
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
5e018de4e5
Quantize module to QQLinear ( #3106 )
2026-02-09 14:35:17 -08:00
ca14d3d835
Fix sharding of quantized models with non-power-of-2 bits ( #3006 )
2026-01-18 07:21:56 -08: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
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
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
4fda5fbdf9
add python testing for cuda with ability to skip list of tests ( #2295 )
2025-06-15 10:56:48 -07:00
c763fe1be0
default strict mode for module update and update_modules ( #2239 )
2025-06-05 15:27:02 -07:00
aa5d84f102
Allow quant layer to be unfrozen ( #2142 )
2025-04-30 09:08:29 -07:00
ca305afdbe
loading empty list is ok when strict = false ( #1834 )
2025-02-05 16:19:27 -08:00
33421c1dd3
Limit grad recursion depth by not recursing through non-grad inputs ( #1764 )
...
* limit grad recursion depth
* add grad of module test
2025-01-14 14:33:18 -08:00
657f466402
use sdpa and exportable functions in transformer multi head attention ( #1760 )
2025-01-09 13:11:55 -08:00
c3628eea49
Add mx.finfo and use it when making causal mask ( #1726 )
...
* finfo
* fixes
* docs
2024-12-19 14:52:41 -08:00
af5a614aad
Eval before cleanup so model file is unlocked ( #1702 )
2024-12-14 21:41:49 -08:00
d8c824c594
Formatting fixes ( #1606 )
2024-11-20 15:30:36 -08:00
cb431dfc9f
Adds 3D pooling ( #1526 )
2024-11-19 16:45:24 -08:00
59247c2b62
add groups in conv2d ( #1569 )
2024-11-07 13:57:53 -08:00
4a64d4bff1
Add support for grouped 1D convolutions to the nn API ( #1444 )
...
* Fix the weight shape for grouped convolutions from the nn API.
* Add tests.
* Pre-commit formatting.
* Add input validation.
* Use integer division instead of casting.
* docs
* nit
---------
Co-authored-by: Awni Hannun <awni@apple.com >
2024-09-28 06:41:07 -07:00
8b30acd7eb
fix module attribute set, reset, set ( #1403 )
2024-09-11 16:30:42 -07:00
a098bc92e0
Fix: Preserve input dtype in Dropout layer output ( #1323 )
...
* Fix: Preserve input dtype in Dropout layer output
- Modified Dropout implementation to ensure that the output dtype matches the input dtype.
- This resolves the issue #1321
* Update test cases in test_nn.py
- Revised test cases to align with updated dropout code
- Fixed assertion method: replaced self.assertTrue with self.assertEqual for accurate comparisons in test_nn.py -> test_rope, test_alibi and test_dropout,
* updated dropout.py
2024-08-13 11:54:21 -07:00
6e06e3a904
feat: Added "tanh" option to GELU approximation ( #1268 )
2024-07-28 09:07:56 +02:00
0fe6895893
Fix the hard-shrink test ( #1185 )
2024-06-04 16:22:56 -07:00
0b7d71fd2f
Add softmin, hardshrink, hardtanh ( #1180 )
...
---------
Co-authored-by: Nikhil Mehta <nikmehta@tesla.com >
2024-06-04 15:48:18 -07:00
84d61d27aa
Make sure 0 is represented in the quantization ( #1016 )
2024-04-19 19:47:26 -07:00
107ba2891a
gelu tanh approx ( #989 )
...
* gelu tanh approx
* gelu tanh approx
* replace gelu approx with tanh approach
* fix comments
* fix comment
2024-04-15 19:49:00 -07:00
cd9e184529
Quantize embedding ( #994 )
...
* quantize embedding
* rename as_linear + comment
* consistency in docs
* fix test
2024-04-15 16:42:10 -07:00
d611251502
Support Chaining for some of functionalities of nn.Module ( #885 ) ( #897 )
...
* add chaining support for some of the functionalities of "nn.Module"
* reformat
* change the return types
* remove return types
* add return type with forward referencing
* add tests for chaining
* add name to contributors
* Update python/mlx/nn/layers/base.py
Co-authored-by: Awni Hannun <awni.hannun@gmail.com >
* Update python/mlx/nn/layers/base.py
Co-authored-by: Awni Hannun <awni.hannun@gmail.com >
* update docstring
* update docstrings
---------
Co-authored-by: Awni Hannun <awni.hannun@gmail.com >
2024-03-27 19:58:29 -07:00
570f2bf29e
pick up preivously set attributes ( #905 )
2024-03-26 11:19:59 -07:00
479051ce1c
add numeric type hierarchy and issubdtype as well as a set_dtype meth… ( #427 )
...
* add numeric type hierarchy and issubdtype as well as a set_dtype method to nn.Module with predicate
numeric type hierarchy and issubtype is compatible to the [numpy hierarchy](https://github.com/numpy/numpy/blob/220f0ab2c5958f9d49b9a28c69a760657b06d9b9/numpy/_core/numerictypes.py#L42 ).
Closes #285 .
* nits in docs
* unify type category checking
* nits in docs
* nits in docs
* more docs nits
* fix callable type
---------
Co-authored-by: Awni Hannun <awni@apple.com >
2024-03-25 12:32:59 -07:00
366478c560
fix modules with dict ( #819 )
2024-03-12 08:54:06 -07:00
8e5600022a
Implement RNN, GRU, LSTM ( #268 )
...
* RNN base implementation
* Address comments+format
* nits in docs
* add tests for prb
* fix test
* add a couple tests
---------
Co-authored-by: Awni Hannun <awni@apple.com >
2024-03-11 21:14:44 -07:00
6a665ea6ed
Dilation for convolutional layers ( #766 )
...
* add dilation parameter to Conv1d layer
* space here too
* add conv1d dilation test
* add dilation parameter for Conv2d layer
* conv2d dilation test
2024-03-04 06:43:00 -08:00
de3d2467a3
Update: Fast GeLU Approximation ( #744 )
...
* add: fast gelu approx
* fix docs
* Update gelu_fast_approx function documentation
* Update python/mlx/nn/layers/activations.py
Co-authored-by: Awni Hannun <awni.hannun@gmail.com >
* fix: test gelu
---------
Co-authored-by: Awni Hannun <awni.hannun@gmail.com >
2024-02-26 21:08:50 -08:00
22364c40b7
Upsample2d ( #414 )
...
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com >
Co-authored-by: Awni Hannun <awni.hannun@gmail.com >
2024-02-23 09:55:04 -08:00
e54cbb7ba6
Pooling layers ( #357 )
...
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com >
Co-authored-by: Awni Hannun <awni@apple.com >
2024-02-12 22:08:13 -08:00
1b97b2958b
Compile with capture ( #629 )
...
* Simple kernel generation
* Remove the generate kernel from graph_utils
* fix multi-output with compile
* fuse with stopgrad
* v1 input, output capture in compile
* cleanup tree update with visitor update
* nit
* remove todo
* state for model, optional explicit init and more pure optimizer steps
* move learning rate to state
* add lr to opt state, some fixes in capture
* fix optim
* update tuple of containers as well
* fix stream for compiled output
* rng state for compile
* nit
* updates and comments
---------
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com >
2024-02-07 17:29:22 -08:00
e5e816a5ef
fix sequential with empty modules at end ( #647 )
2024-02-07 13:22:27 -08:00
09b9275027
Make shape a tuple ( #591 )
...
* shape tuple
* also remove simplify from docs
* rebase
2024-01-30 13:11:01 -08:00
d3a9005454
Softshrink mapping + op ( #552 )
...
* Added Softshrink mapping + op
* formatting
* docs + nits in docstring
---------
Co-authored-by: Awni Hannun <awni@apple.com >
2024-01-30 12:56:28 -08:00
f6feb61f92
feat: add support for saving safetensors in the save_weights ( #497 )
...
* feat: add save safetensors support in module save_weights
* chore: checking missing changes
* Update python/mlx/nn/layers/base.py
Co-authored-by: Awni Hannun <awni.hannun@gmail.com >
* chore: update docstring for load_weights
---------
Co-authored-by: Awni Hannun <awni.hannun@gmail.com >
2024-01-19 06:19:33 -08:00
e9ca65c939
Fix BN stats to not expand shape ( #409 )
...
* fix BN stats to not expand shape
* nit
2024-01-09 11:54:51 -08:00
022a944367
Added GLU activation function and Gated activation function ( #329 )
...
* Added GLU activation function and gated activation function
* Ran pre-commit
* Ran pre commit
* Removed old sigmoid implementation to match with main
* Removed gated activation from __init__.py
* Removed unused test cases
* Removed unused imports
* format / docstring
---------
Co-authored-by: Awni Hannun <awni@apple.com >
2024-01-08 06:13:16 -08:00
75dc537e44
Fix the sigmoid module ( #371 )
2024-01-04 13:16:36 -08:00
cf88db44b5
revert copy ( #366 )
2024-01-04 10:43:29 -08:00
d2467c320d
Added support for python copy ( #335 )
...
* Added support for python copy
* precommit changes
* removed `_compiled_call_impl` line
* added tests and suggested changes
* ACK changes
2024-01-03 20:59:40 -08:00
Shantanu Suryawanshi