jax-ml/jax · error · NotImplementedError

{uninitialized_memory} + {str(dtype)}

Error message

{uninitialized_memory} + {str(dtype)}

What it means

get_uninitialized_value could not produce an uninitialized sentinel for the combination of uninitialized_memory mode and dtype. Only NaN (floating) and zero modes are supported; NaN for non-float or unrecognized dtype combos hits the NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic/interpret/utils.py:43

from jax._src.pallas import primitives
from jax._src.util import safe_map
import jax.numpy as jnp
import numpy as np


def get_uninitialized_value(
    dtype, uninitialized_memory: Literal["nan", "zero"]
):
  if uninitialized_memory == "nan":
    if jnp.issubdtype(dtype, jnp.floating):
      return np.nan
    elif jnp.issubdtype(dtype, jnp.integer):
      return jnp.iinfo(dtype).max
    elif jnp.issubdtype(dtype, jnp.bool):
      return True
  if uninitialized_memory == "zero":
    return 0
  raise NotImplementedError(uninitialized_memory + " + " + str(dtype))


def get_uninitialized_array(
    shape, dtype, uninitialized_memory: Literal["nan", "zero"]
):
  return jnp.full(
      shape,
      get_uninitialized_value(dtype, uninitialized_memory),
      dtype,
  )


def pad_to_block_dimension(
    value, block_shape, uninitialized_memory: Literal["nan", "zero"]
):
  """Pads values so the shape evenly divides into block dimensions.

  For example, if values has a shape of (33, 2, 5) with a block_shape of

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use uninitialized_memory='zero' if the dtype has no NaN sentinel
2. Avoid depending on OOB/padded reads for non-float dtypes; mask reads instead
3. File/patch support for the dtype in get_uninitialized_value if it should be supported

Example fix

// before
params = interpret.InterpretParams(uninitialized_memory='nan')  # with int32 buffers
// after
params = interpret.InterpretParams(uninitialized_memory='zero')

import jax.numpy as jnp
mode = 'nan' if jnp.issubdtype(dtype, jnp.floating) else 'zero'

try:
    val = get_uninitialized_value(dtype, 'nan')
except NotImplementedError:
    val = 0  # fall back to zero for non-float dtypes

When it happens

Trigger: Requesting uninitialized (padded/out-of-bounds) memory with uninitialized_memory='nan' for an integer/bool/bfloat16-unsupported dtype, or an unrecognized dtype, in interpret-mode reads.

Common situations: Kernels reading padded regions of integer accumulators in interpret mode; new/extended dtypes not yet covered by the sentinel table.

Related errors

• bitcast_convert_type with different bitwidths not supported
• masked load_p
• run_scoped_p with collective axes is not supported
• Non-decrementing wait is not supported.
• Vector clock size ({self.vector_clock_size}) must be greater

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/d7ea9f7e7441c44f. Report an issue: GitHub.