{"record":{"id":"cb48a1a4f0956ff9","repo":"jax-ml/jax","slug":"the-is-fully-replicated-property-was-called-on-se","errorCode":null,"errorMessage":"The is_fully_replicated property was called on {self._error_repr()}.{self._origin_msg()}","messagePattern":"The is_fully_replicated property was called on (.+?)\\.(.+?)","errorType":"exception","errorClass":"ConcretizationTypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":1267,"sourceCode":"  def global_shards(self):\n    raise ConcretizationTypeError(self,\n      f\"The global_shards property was called on {self._error_repr()}.\"\n      f\"{self._origin_msg()}\")\n\n  def is_deleted(self):\n    raise ConcretizationTypeError(self,\n      f\"The is_deleted() method was called on {self._error_repr()}.\"\n      f\"{self._origin_msg()}\")\n\n  @property\n  def is_fully_addressable(self):\n    raise ConcretizationTypeError(self,\n      f\"The is_fully_addressable property was called on {self._error_repr()}.\"\n      f\"{self._origin_msg()}\")\n\n  @property\n  def is_fully_replicated(self):\n    raise ConcretizationTypeError(self,\n      f\"The is_fully_replicated property was called on {self._error_repr()}.\"\n      f\"{self._origin_msg()}\")\n\n  def on_device_size_in_bytes(self):\n    raise ConcretizationTypeError(self,\n      f\"The on_device_size_in_bytes() method was called on {self._error_repr()}.\"\n      f\"{self._origin_msg()}\")\n\n  @property\n  def traceback(self):\n    raise ConcretizationTypeError(self,\n      f\"The traceback property was called on {self._error_repr()}.\"\n      f\"{self._origin_msg()}\")\n\n  def unsafe_buffer_pointer(self):\n    raise ConcretizationTypeError(self,\n      f\"The unsafe_buffer_pointer() method was called on {self._error_repr()}.\"\n      f\"{self._origin_msg()}\")","sourceCodeStart":1249,"sourceCodeEnd":1285,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L1249-L1285","documentation":"ConcretizationTypeError raised when the Tracer.is_fully_replicated property is accessed. is_fully_replicated reports whether all shards of a materialized jax.Array hold identical data (relevant to sharded/distributed arrays); tracers carry no layout, so the property stub raises during tracing.","triggerScenarios":"Accessing `x.is_fully_replicated` on a traced value inside jit/pjit/grad/vmap/scan/remat or inside a custom transform rule.","commonSituations":"Optimization branches like `if x.is_fully_replicated: skip_all_reduce(x)` placed inside a jitted collective function; distributed training code refactored under vmap; sharding assertions copied into traced kernels.","solutions":["Move the replication check outside the traced function and pass the result as a static boolean.","Derive replication info from the value's sharding on its aval (e.g. via jax.typeof) where tracing-time knowledge is required.","Let the XLA/SPMD compiler handle replication-aware elision of collectives rather than branching in Python.","Guard with isinstance(x, jax.Array) in shared helper code."],"exampleFix":"// before\n@jax.jit\ndef allreduce(x):\n    if x.is_fully_replicated:  # ConcretizationTypeError\n        return x\n    return lax.psum(x, 'i')\n\n// after\ndef allreduce(x):\n    if x.is_fully_replicated:   # concrete check outside\n        return x\n    return _psum_jit(x)","handlingStrategy":"validation","validationCode":"import jax\n\ndef fully_replicated_or_default(x, default=False):\n    if isinstance(x, jax.core.Tracer):\n        return default\n    return x.is_fully_replicated","typeGuard":"import jax\nfrom jax.core import Tracer\n\ndef replication_known(x) -> bool:\n    return isinstance(x, jax.Array) and not isinstance(x, Tracer)","tryCatchPattern":"from jax.errors import ConcretizationTypeError\ntry:\n    rep = x.is_fully_replicated\nexcept ConcretizationTypeError:\n    rep = False  # conservative fallback under tracing","preventionTips":["Hoist replication checks to the caller; pass as static boolean.","Let the compiler elide redundant collectives instead of Python branching.","Document which helper functions are trace-safe vs host-only."],"tags":["jax","concretization","tracer","sharding","replication"],"backgroundTag":"jax-tracer-concretization-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}