code: def closure_to_pytree(tree):
      """Convert all function closures into pytree nodes.
  
      **Arguments:**
  
      - `tree`: Any pytree.
  
      **Returns:**
  
      A copy of `tree`, where all function closures have been replaced by a new object
      that is (a) callable like the original function, but (b) iterates over its
      `__closure__` as subnodes in the pytree.
  
      !!! Example
  
          ```python
          def some_fn():
              a = jnp.array(1.)
  
              @closure_to_pytree
              def f(x):
                  return x + a
  
              print(jax.tree_util.tree_leaves(f))  # prints out `a`
          ```
  
      !!! Warning
  
          One annoying technical detail in the above example: we had to wrap the whole lot
          in a `some_fn`, so that we're in a local scope. Python treats functions at the
          global scope differently, and this conversion won't result in any global
          variable being treated as part of the pytree.
  
          In practice, the intended use case of this function is to fix Optax, which
          always uses local functions.
      """
      return jtu.tree_map(_fixup_closure, tree)
  
file_path: equinox/equinox/internal/_closure_to_pytree.py
imports: import inspect
  import types
  from typing import Any, Optional
  import jax.tree_util
  from .._module import Module
label_desc: Use image node_id 4 for calling a global function with example usage: closure_to_pytree(tree) and returns: jtu