Awni Hannun
1b97b2958b
* 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>
38 lines
677 B
ReStructuredText
38 lines
677 B
ReStructuredText
Module
|
|
======
|
|
|
|
.. currentmodule:: mlx.nn
|
|
|
|
.. autoclass:: Module
|
|
|
|
.. rubric:: Attributes
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
|
|
Module.training
|
|
Module.state
|
|
|
|
.. rubric:: Methods
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
|
|
Module.apply
|
|
Module.apply_to_modules
|
|
Module.children
|
|
Module.eval
|
|
Module.filter_and_map
|
|
Module.freeze
|
|
Module.leaf_modules
|
|
Module.load_weights
|
|
Module.modules
|
|
Module.named_modules
|
|
Module.parameters
|
|
Module.save_weights
|
|
Module.train
|
|
Module.trainable_parameters
|
|
Module.unfreeze
|
|
Module.update
|
|
Module.update_modules
|