Fix compile with non standard types (#745)

* refactor tree utils

* fix compile + tree code refactor

* Add an extra test

* add a few missing activations to docs

* hash structure

* Encode the full argument structure

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

python/mlx/nn/layers/__init__.py

@@ -37,6 +37,7 @@ from mlx.nn.layers.activations import (
     relu,
     relu6,
     selu,
+    sigmoid,
     silu,
     softmax,
     softplus,

python/mlx/nn/layers/activations.py

@@ -18,7 +18,7 @@ def _make_activation_module(f):
 
 @partial(mx.compile, shapeless=True)
 def sigmoid(x):
-    r"""Applies the element-wise function:
+    r"""Applies the sigmoid function.
 
     .. math::
         \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)}