code: def test_grad_of_triple_nested_for_loop(self):
      func = lambda x: jnp.sin(x) + 1.0
  
      @jax.jit
      def f(x):
          out = jnp.zeros_like(x)
  
          def body(i, j, k, refs):
              x_ref, out_ref = refs
              y = func(x_ref[i, j, k])
              out_ref[i, j, k] += y
  
          return for_loop.for_loop(x.shape, body, (x, out))[1].sum()
  
      x = random.normal(random.PRNGKey(0), (5, 4, 3))
      ref = lambda x: jax.vmap(jax.vmap(jax.vmap(func)))(x).sum()
      self.assertAllClose(f(x), ref(x))
      jtu.check_grads(f, (x,), order=2, atol=0.1, rtol=0.1)
  
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 6 for calling the ForLoopTransformationTest obj's underlying member method code with example usage: obj.test_grad_of_triple_nested_for_loop() and returns: for_loop