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>

docs/src/python/optimizers.rst

@@ -29,14 +29,16 @@ model's parameters and the **optimizer state**.
             # Compute the new parameters but also the optimizer state.
             mx.eval(model.parameters(), optimizer.state)
 
+.. toctree::
+
+   optimizer
+
 .. currentmodule:: mlx.optimizers
 
 .. autosummary::
    :toctree: _autosummary
    :template: optimizers-template.rst
 
-   OptimizerState
-   Optimizer
    SGD
    RMSprop
    Adagrad