code: def test_for_loop_invar(self):
      def f(x):
          s = jnp.ones((2, 32), x.dtype)
  
          def body(i, refs):
              x_ref, y_ref = refs
              y_ref[i] = s * x_ref[i] * jnp.cos(s)
              # We should save `s` and `jnp.cos(s)` as residuals and not broadcast
              # them.
  
          return for_loop.for_loop(
              x.shape[0], body, (x, jnp.zeros_like(x))
          )
  
      _, f_vjp = jax.linearize(f, jnp.ones((5, 2, 32)))
      jaxpr = jax.make_jaxpr(f_vjp)(jnp.ones((5, 2, 32)))
      consts = [
          v.aval
          for v in jaxpr.jaxpr.constvars
          if v.aval.shape == (2, 32)
      ]
      self.assertLen(consts, 2)
  
      def loss(A):
          def step(x, _):
              return jnp.matmul(A, x), None
  
          init_x = jnp.zeros(A.shape[-1:])
          last_x, _ = for_loop.scan(step, init_x, jnp.arange(10))
          return jnp.sum(last_x)
  
      A = jnp.zeros((3, 3))
      # The second DUS was unnecessarily replicating A across time.
      # We check XLA because _scan_impl is "underneath" the jaxpr language.
      s = str(jax.xla_computation(jax.grad(loss))(A).as_hlo_text())
      assert s.count("dynamic-update-slice(") < 2
  
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 3 for calling the ForLoopTransformationTest obj's underlying member method code with example usage: obj.test_for_loop_invar() and returns: for_loop, jnp, jnp,  None