code: def test_onnx_export():
      @jax.vmap
      def fn(x, y):
          x = x + 1
          y = eqxi.unvmap_any(y)
          return jnp.where(y, x, 1)
  
      onnx_fn = eqxi.to_onnx(fn)
      args = jnp.array([1, 2]), jnp.array([True, False])
      onnx_fn(*args)
  
imports: import equinox.internal
  import jax
  import jax.numpy
  import pytest
label_desc: Use image node_id 1 for calling a global function with example usage: test_onnx_export() and returns: jnp