code: def test_param_with_axes_no_axes(self, axes_spec):
      class ParamTest(nn.Module):
          @nn.compact
          def __call__(self, x):
              foo = partitioning.param_with_axes(
                  "foo",
                  lambda k, s, d: jnp.zeros(s, d),
                  (2, 2),
                  x.dtype,
                  axes=axes_spec,
              )
              return x + foo
  
      k = random.key(0)
      x = jnp.ones((2, 2))
      _ = ParamTest().init(k, x)
  
imports: import jax
  import jax.numpy
  from absl.testing import absltest, parameterized
  from jax import random, sharding
  from jax.experimental import mesh_utils
  from flax import linen
  from flax.core import freeze, unfreeze
  from flax.linen import partitioning
label_desc: Use image node_id 8 for calling the PartitioningTest obj's underlying member method code with example usage: obj.test_param_with_axes_no_axes(axes_spec) and returns: unknown