code: def test_nondifferentiable_body1():
      def cond_fun(carry):
          return True
  
      def body_fun(carry):
          step, x, y, z, _ = carry
          y2 = eqxi.nondifferentiable(y)
          return step + 1, x + y2, y + 1, z.at[step].set(y), None
  
      @eqx.filter_jit
      @jax.value_and_grad
      def run(x__z, y_in, true):
          x_in, z_in = x__z
          init = (0, x_in, y_in, z_in, None)
          if true:
              out = _while_as_scan(
                  cond_fun, body_fun, init, max_steps=3
              )
          else:
              out = eqxi.while_loop(
                  cond_fun,
                  body_fun,
                  init,
                  max_steps=3,
                  kind="checkpointed",
              )
          _, x_out, y_out, z_out, none = out
          assert none is None
          return x_out + y_out + jnp.sum(z_out)
  
      x_in = jnp.array(1.2)
      y_in = jnp.array(0.7)
      z_in = jnp.array([-5.0, -5.0, -5.0])
      true = run((x_in, z_in), y_in, true=True)
      outs = run((x_in, z_in), y_in, true=False)
      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 19 for calling a global function with example usage: test_nondifferentiable_body1() and returns: True, unknown,  unknown,  unknown,  z,  None, unknown