Merge pull request #4374 from 8bitmp3:add-nnx-eval_shap-docstring

PiperOrigin-RevId: 696587888

flax/nnx/transforms/transforms.py

@@ -138,7 +138,14 @@ def _eval_shape_fn(*args, **kwargs):
 
   out = jax.eval_shape(_eval_shape_fn, *args, **kwargs)
   return extract.from_tree(out)
-
+  """A "lifted" version of `jax.eval_shape <https://jax.readthedocs.io/en/latest/_autosummary/jax.eval_shape.html#jax.eval_shape>`_
+    that can handle `flax.nnx.Module <https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html#flax.nnx.Module>`_
+    / graph nodes as arguments.
+
+  Similar to ``jax.eval_shape``, it computes the shape/dtype of a function `f` without
+    performing any floating point operations (FLOPs) which can be expensive. This can be
+    useful for performing shape inference, for example.
+  """
 
 # -------------------------------
 # cond and switch