code: def test_complicated(getkey):
      num_lowerings = 0
  
      def increment(stack):
          nonlocal num_lowerings
          if stack.endswith("mlir"):
              num_lowerings += 1
  
      def call(f, x):
          return f(x + 1)
  
      @eqx.internal.noinline
      def call_noinline(f, x):
          print("hi")
          x = eqx.internal.announce_transform(x, announce=increment)
          return call(f, x)
  
      mlp = eqx.nn.MLP(1, 1, 16, 2, key=getkey())
  
      @jax.jit
      @jax.vmap
      @jax.grad
      def run(x):
          x = x[None]
          y = call(mlp, x)
          z = call(mlp, y)
          w = call(mlp, z)
          (out,) = w * x
          return out
  
      @jax.jit
      @jax.vmap
      @jax.grad
      def run_noinline(x):
          x = x[None]
          y = call_noinline(mlp, x)
          z = call_noinline(mlp, y)
          w = call_noinline(mlp, z)
          (out,) = w * x
          return out
  
      # Note that the w*x in the above is actually meaningful.
      # This promotes the cotangent on the loss (=1.0) to something
      # with a batch dimension.
      # Otherwise the `w = call_noinline(mlp, z)` receives a cotangent
      # without a batch dimension, whilst all the others do get a batch
      # dimension, and that triggers an additional compilation.
  
      xs = jnp.array([1.0, 2.0, 3.0])
  
      assert tree_allclose(run(xs), run_noinline(xs))
      assert num_lowerings == 3
  
file_path: equinox/tests/test_noinline.py
imports: import equinox
  import jax
  import jax.numpy
  import jax.random
  from .helpers import tree_allclose
label_desc: Use image node_id 9 for calling a global function with example usage: test_complicated(getkey) and returns: f, call, out, out