code: def test_fixup_optax():
      lr = jnp.array(1e-3)
      optim = optax.chain(
          optax.adam(lr),
          optax.scale_by_schedule(
              optax.piecewise_constant_schedule(1, {200: 0.1})
          ),
      )
      optim = eqxi.closure_to_pytree(optim)
  
      for leaf in jtu.tree_leaves(optim):
          if eqx.is_array(leaf) and leaf == -lr:
              break
      else:
          assert False
  
      # Check that we can still init and update as normal.
      grads = params = {"foo": jnp.array(1.0)}
      state = optim.init(params)
      with jax.numpy_dtype_promotion("standard"):
          optim.update(grads, state)
  
      lr = jnp.array(1e-2)
      optim2 = optax.chain(
          optax.adam(lr),
          optax.scale_by_schedule(
              optax.piecewise_constant_schedule(1, {200: 0.1})
          ),
      )
      optim2 = eqxi.closure_to_pytree(optim2)
  
      compiling = 0
  
      @eqx.filter_jit
      def f(x):
          nonlocal compiling
          compiling += 1
  
      f(optim)
      assert compiling == 1
      f(optim2)
      assert compiling == 1
  
file_path: equinox/tests/test_closure_to_pytree.py
imports: import equinox
  import equinox.internal
  import jax
  import jax.numpy
  import jax.tree_util
  import optax
label_desc: Use image node_id 1 for calling a global function with example usage: test_fixup_optax() without return types