code: def test_with_sharding_constraint_fallback(self, wsc_fn):
      arr = jnp.ones((2, 2))
      with partitioning.axis_rules(AXIS_RULES_1):
          _ = partitioning.with_sharding_constraint(
              arr, ("foo", "not_recognized")
          )
          wsc_fn.assert_called_with(
              arr, jax.sharding.PartitionSpec("data", None), mesh=None
          )
          wsc_fn.reset_mock()
          _ = partitioning.with_sharding_constraint(
              arr,
              ("foo", "not_recognized"),
              fallback=partitioning.RulesFallback.AXIS_IS_UNSHARDED,
          )
          wsc_fn.assert_called_with(
              arr, jax.sharding.PartitionSpec("data", None), mesh=None
          )
          wsc_fn.reset_mock()
          with self.assertRaises(ValueError):
              _ = partitioning.with_sharding_constraint(
                  arr,
                  ("foo", "not_recognized"),
                  fallback=partitioning.RulesFallback.RAISE_ERROR,
              )
          wsc_fn.assert_not_called()
          _ = partitioning.with_sharding_constraint(
              arr,
              ("foo", "not_recognized"),
              fallback=partitioning.RulesFallback.NO_CONSTRAINT,
          )
          wsc_fn.assert_not_called()
  
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 7 for calling the PartitioningTest obj's underlying member method code with example usage: obj.test_with_sharding_constraint_fallback(wsc_fn) without return types