code: def test_body_fun_grads(capfd):
      def cond_fun(carry):
          return True
  
      @eqx.filter_jit
      @jax.grad
      def run(x__y, true):
          x, y = x__y
          # `init` and `body_fun` are deliberately chosen so that `carry[0]` requires a
          # gradient solely for the purpose of propagating that gradient back into `x`.
          # (And in particular, not for propagating it back into `init`.)
          # Thus this test is checking that we get gradients with respect to `body_fun`
          # correctly.
          init = (1.0, y)
  
          def body_fun(carry):
              a, b = carry
              return a * x, b + 1
  
          if true:
              final = _while_as_scan(
                  cond_fun, body_fun, init, max_steps=3
              )
          else:
              final = eqxi.while_loop(
                  cond_fun,
                  body_fun,
                  init,
                  max_steps=3,
                  kind="checkpointed",
              )
          return sum(final)
  
      x__y = (jnp.array(1.0), jnp.array(1.0))
  
      capfd.readouterr()
      outs = run(x__y, true=False)
      text, _ = capfd.readouterr()
      assert "perturb_val (False, False, False, (True, True))" in text
      assert (
          "symbolic_zero_gradient (True, True, True, (False, False))"
          in text
      )
  
      true = run(x__y, true=True)
      assert tree_allclose(true, outs)
  
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 21 for calling a global function with example usage: test_body_fun_grads(capfd) and returns: True, sum, unknown,  unknown