code: def test_lift_with_vmap(self):
      def inner_fn(x):
          assert x.ndim == 1
          return Bias()(x)
  
      def outer_fn(x):
          assert x.ndim == 2
          x = Bias()(x)
          inner = multi_transform.without_apply_rng(
              transform.transform(inner_fn)
          )
          inner_p = lift.lift(inner.init)(base.next_rng_key(), x[0])
          vmap_inner = jax.vmap(inner.apply, in_axes=(None, 0))
          return vmap_inner(inner_p, x)
  
      key = jax.random.PRNGKey(428)
      init_key, apply_key = jax.random.split(key)
      data = np.zeros((3, 2))
  
      outer = transform.transform(outer_fn)
      outer_params = outer.init(init_key, data)
      self.assertEqual(
          outer_params,
          {
              "bias": {"b": np.ones(())},
              "lifted/bias": {"b": np.ones(())},
          },
      )
  
      out = outer.apply(outer_params, apply_key, data)
      np.testing.assert_equal(out, 2 * np.ones((3, 2)))
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import base
  from haiku._src import config
  from haiku._src import lift
  from haiku._src import module
  from haiku._src import multi_transform
  from haiku._src import stateful
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 3 for calling the LiftTest obj's underlying member method code with example usage: obj.test_lift_with_vmap() and returns: Bias, vmap_inner