{"record":{"id":"1fce5943967cd5e6","repo":"jax-ml/jax","slug":"the-on-device-size-in-bytes-method-was-called-on","errorCode":null,"errorMessage":"The on_device_size_in_bytes() method was called on {self._error_repr()}.{self._origin_msg()}","messagePattern":"The on_device_size_in_bytes\\(\\) method was called on (.+?)\\.(.+?)","errorType":"exception","errorClass":"ConcretizationTypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":1272,"sourceCode":"  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()}\")\n\n_jax.set_tracer_class(Tracer)\n\n# these can be used to set up forwarding of properties and instance methods from\n# Tracer instances to the underlying avals","sourceCodeStart":1254,"sourceCodeEnd":1290,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L1254-L1290","documentation":"ConcretizationTypeError raised when Tracer.on_device_size_in_bytes() is called. This method reports the device-side byte size of a materialized jax.Array's buffers; a tracer has no storage, so the base-class stub raises during any JAX transform's tracing.","triggerScenarios":"Calling `x.on_device_size_in_bytes()` on a value traced by jit/grad/vmap/pmap/scan/remat or a custom rule.","commonSituations":"Memory-budget accounting or telemetry (`total += x.on_device_size_in_bytes()`) wired into a jitted step; profiling helpers reused under vmap; OOM-debugging instrumentation left in traced code.","solutions":["Compute size accounting outside the trace, on the concrete arrays passed to/returned from jit.","Estimate sizes statically from shapes/dtypes (x.aval.shape, itemsize) if you need them during tracing.","Remove instrumentation from traced kernels and collect it in the surrounding Python driver.","Cache per-shape byte sizes keyed by input shapes to avoid repeated host queries."],"exampleFix":"// before\n@jax.jit\ndef step(x):\n    log(x.on_device_size_in_bytes())  # ConcretizationTypeError\n    return heavy(x)\n\n// after\ndef step_outer(x):\n    log(x.on_device_size_in_bytes())  # concrete array\n    return _step_jit(x)","handlingStrategy":"type-guard","validationCode":"import jax\n\ndef device_bytes(x):\n    if isinstance(x, jax.core.Tracer):\n        import numpy as np\n        a = x.aval\n        return int(np.prod(a.shape)) * a.dtype.itemsize  # static estimate\n    return x.on_device_size_in_bytes()","typeGuard":"import jax\nfrom jax.core import Tracer\n\ndef size_queryable(x) -> bool:\n    return isinstance(x, jax.Array) and not isinstance(x, Tracer)","tryCatchPattern":"from jax.errors import ConcretizationTypeError\ntry:\n    nbytes = x.on_device_size_in_bytes()\nexcept ConcretizationTypeError:\n    nbytes = None  # use static shape-based estimate instead","preventionTips":["Do memory accounting in the host driver on concrete arrays.","Estimate sizes from shape/dtype when inside transforms.","Keep telemetry calls out of jitted kernels."],"tags":["jax","concretization","tracer","memory","profiling"],"backgroundTag":"jax-tracer-concretization-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}