code: def _with_sharding_constraint(
      x: Array,
      axis_resources: tp.Optional[jax.sharding.PartitionSpec],
      mesh: tp.Optional[jax.sharding.Mesh] = None,
  ):
      # if jax.devices()[0].platform == "cpu" or (
      if not _global_mesh_defined() and mesh is None:
          return x
      else:
          if mesh is not None and axis_resources is not None:
              sharding = jax.sharding.NamedSharding(
                  mesh, axis_resources
              )
              return jax.lax.with_sharding_constraint(x, sharding)
          return jax.lax.with_sharding_constraint(x, axis_resources)
  
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 7 for calling a global function with example usage: _with_sharding_constraint(x, axis_resources, mesh) and returns: x, jax, jax