code: def nested_for_loop(nsteps, body, state, **kwargs):
      def outer_body(_, refs):
          def inner_body(i, _):
              body(i, refs)
              return
  
          for_loop.for_loop(nsteps, inner_body, ())
  
      return for_loop.for_loop(1, outer_body, state)
  
imports: from functools import partial
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  import jax
  from jax import random
  from jax._src import test_util
  from jax._src.lax.control_flow import for_loop
  import jax.numpy
  from jax import config
label_desc: Use image node_id 2 for calling a global function with example usage: nested_for_loop(nsteps, body, state) and returns: for_loop