code: def _get_problem(key, *, num_steps: Optional[int]):
      valkey1, valkey2, modelkey = jr.split(key, 3)
  
      def cond_fun(carry):
          if num_steps is None:
              return True
          else:
              step, _, _ = carry
              return step < num_steps
  
      def make_body_fun(dynamic_mlp):
          mlp = eqx.combine(dynamic_mlp, static_mlp)
  
          def body_fun(carry):
              # A simple new_val = mlp(val) tends to converge to a fixed point in just a
              # few iterations, which implies zero gradient... which doesn't make for a
              # test that actually tests anything. Making things rotational like this
              # keeps things more interesting.
              step, val1, val2 = carry
              (theta,) = mlp(val1)
              real, imag = val1
              z = (
                  real.astype(jnp.complex64)
                  + imag.astype(jnp.complex64) * 1j
              )
              z = z * jnp.exp(1j * theta.astype(jnp.complex64))
              real = jnp.real(z)
              imag = jnp.imag(z)
              jax.debug.print("{}", step)  # pyright: ignore
              val1 = jnp.stack([real, imag])
              val2 = val2.at[step].set(real)
              return step + 1, val1, val2
  
          return body_fun
  
      init_val1 = jr.normal(valkey1, (2,))
      init_val2 = jr.normal(valkey2, (20,))
      if num_steps is not None:
          # So that things fit in the buffer update above.
          assert num_steps < 20
      mlp = eqx.nn.MLP(2, 1, 2, 2, key=modelkey)
      dynamic_mlp, static_mlp = eqx.partition(mlp, eqx.is_array)
  
      return cond_fun, make_body_fun, init_val1, init_val2, dynamic_mlp
  
imports: import functools
  import timeit
  from typing import Optional
  import equinox
  import equinox.internal
  import jax
  import jax.lax
  import jax.lib
  import jax.numpy
  import jax.random
  import jax.tree_util
  import pytest
  from .helpers import tree_allclose
label_desc: Use image node_id 1 for calling a global function with example usage: _get_problem(key) and returns: cond_fun,  make_body_fun,  init_val1,  init_val2,  dynamic_mlp, body_fun, True, unknown, unknown,  val1,  val2