code: def _with_sharding_constraint(
      x: Array,
      axis_resources: Optional[jax.sharding.PartitionSpec],
      mesh: Optional[jax.sharding.Mesh] = None,
  ):
      """Wrapper for lax.with_sharding_constraint, no-op on cpu or outside jit."""
      if jax.devices()[0].platform == "cpu" or (
          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 lax.with_sharding_constraint(x, sharding)
          return lax.with_sharding_constraint(x, axis_resources)
  
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 10 for calling a global function with example usage: _with_sharding_constraint(x, axis_resources, mesh) and returns: x, lax, lax