code: def with_sharding_constraint(
      x: ArrayPytree,
      axis_resources: PartitionSpecPytree,
      mesh: tp.Optional[jax.sharding.Mesh] = None,
  ):
      # If no axis binding is set, this is a no-op.
      if axis_resources is None:
          return x
      # Translate logical names to mesh assignments.
      return jax.tree_util.tree_map(
          functools.partial(_with_sharding_constraint, mesh=mesh),
          x,
          axis_resources,
          is_leaf=_is_spec,
      )
  
imports: import functools
  import typing
  import jax
  from jax.experimental import maps
  from jax.sharding import Mesh, PartitionSpec
  from flax.experimental.nnx.nnx import variables
  from flax.experimental.nnx.nnx.pytreelib import TreeNode
  from flax.experimental.nnx.nnx.state import State
label_desc: Use image node_id 9 for calling a global function with example usage: with_sharding_constraint(x, axis_resources, mesh) and returns: jax, x