code: def test_vmap_unshared(self):
      x = random.normal(random.key(0), (1, 4))
      x = jnp.concatenate([x, x], 0)
  
      y, variables = init(mlp_vmap)(
          random.key(1), x, share_params=False
      )
  
      param_shapes = unfreeze(
          jax.tree_util.tree_map(jnp.shape, variables["params"])
      )
      self.assertEqual(
          param_shapes,
          {
              "hidden_0": {"kernel": (2, 4, 8), "bias": (2, 8)},
              "out": {"kernel": (2, 8, 1), "bias": (2, 1)},
          },
      )
      self.assertEqual(y.shape, (2, 1))
      self.assertFalse(jnp.allclose(y[0], y[1]))
  
file_path: flax/tests/core/design/core_vmap_test.py
imports: from typing import Callable, Sequence
  import jax
  from absl.testing import absltest
  from jax import numpy
  from jax import random
  from flax.core import Array, Scope, init, lift, nn, unfreeze
label_desc: Use image node_id 2 for calling the VMapTest obj's underlying member method code with example usage: obj.test_vmap_unshared() without return types