code: def test_with_sharding_constraint(self, wsc_fn):
      arr = jnp.ones((2, 2))
      axes = ("foo", "bar")
      partitioning.set_axis_rules(())
      _ = partitioning.with_sharding_constraint(arr, axes)
      wsc_fn.assert_not_called()
      with partitioning.axis_rules(AXIS_RULES_1):
          _ = partitioning.with_sharding_constraint(arr, None)
          wsc_fn.assert_not_called()
          _ = partitioning.with_sharding_constraint(arr, axes)
          wsc_fn.assert_called_with(
              arr,
              jax.sharding.PartitionSpec("data", "model"),
              mesh=None,
          )
  
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 6 for calling the PartitioningTest obj's underlying member method code with example usage: obj.test_with_sharding_constraint(wsc_fn) without return types