code: def test_nested_sequential(inner_stateful, outer_stateful, getkey):
      def make():
          inner_seq = eqx.nn.Sequential(
              [
                  eqx.nn.Linear(2, 4, key=getkey()),
                  eqx.nn.BatchNorm(4, axis_name="batch")
                  if inner_stateful
                  else eqx.nn.Identity(),
                  eqx.nn.Linear(4, 3, key=getkey()),
              ]
          )
          outer_seq = eqx.nn.Sequential(
              [
                  eqx.nn.Linear(5, 2, key=getkey()),
                  inner_seq,
                  eqx.nn.BatchNorm(3, axis_name="batch")
                  if outer_stateful
                  else eqx.nn.Identity(),
                  eqx.nn.Linear(3, 6, key=getkey()),
              ]
          )
          return outer_seq
  
      x = jrandom.normal(getkey(), (1, 5))
      if inner_stateful or outer_stateful:
          model, state = eqx.nn.make_with_state(make)()
          out, state = jax.vmap(
              model,
              axis_name="batch",
              in_axes=(0, None),
              out_axes=(0, None),
          )(x, state)
          assert isinstance(state, eqx.nn.State)
      else:
          model = make()
          out = jax.vmap(model, axis_name="batch")(x)
      assert isinstance(out, jax.Array)
      assert out.shape == (1, 6)
  
imports: import warnings
  from typing import Union
  import equinox
  import jax
  import jax.nn
  import jax.numpy
  import jax.random
  import pytest
label_desc: Use image node_id 10 for calling a global function with example usage: test_nested_sequential(inner_stateful, outer_stateful, getkey) and returns: outer_seq