code: def with_logical_constraint(
      x: ArrayPytree,
      logical_axis_resources: LogicalPartitionSpecPytree,
      rules: Optional[LogicalRules] = None,
      mesh: Optional[jax.sharding.Mesh] = None,
      fallback: RulesFallback = RulesFallback.AXIS_IS_UNSHARDED,
  ):
      """Version of jit's with_sharding_constraint that uses logical axis names."""
      # If no axis binding is set, this is a no-op.
      if rules is None:
          rules = _axis_rules.rules
      if not rules or logical_axis_resources is None:
          return x
      # Translate logical names to mesh assignments.
      return jax.tree_util.tree_map(
          functools.partial(
              _with_sharding_constraint_one_fallback,
              fallback=fallback,
              rules=rules,
              mesh=mesh,
          ),
          logical_axis_resources,
          x,
          is_leaf=_is_logical_spec,
      )
  
imports: import collections
  import contextlib
  import dataclasses
  import enum
  import functools
  import threading
  from typing import Any, Callable, List, Optional, Sequence, Tuple, Union
  import jax
  from jax import lax
  from jax.experimental import maps
  from flax import struct
  from flax.core import meta
label_desc: Use image node_id 13 for calling a global function with example usage: with_logical_constraint(x, logical_axis_resources, rules, mesh, fallback) and returns: jax, x