jax-ml/jax · error · ValueError

numpy masked arrays are not supported as direct inputs to JA

Error message

numpy masked arrays are not supported as direct inputs to JAX functions. Use arr.filled() to convert the value to a standard numpy array.

What it means

jaxlib's C++ PyTree flattening machinery calls a registered custom PyTree node's to_iterable_with_keys hook and requires it to return a 2-element tuple (key_leaf_pairs, aux_data) where the first element is iterable. This error is thrown when the first element cannot be cast to a Python iterable (e.g. it is an int, None, or a non-iterable object).

Source

Thrown at jax/_src/abstract_arrays.py:54

} | {np.dtype(dt).type for dt in dtypes._float_types}

if dtypes.int2 is not None:
  assert dtypes.uint2 is not None
  numpy_scalar_types.add(dtypes.int2)
  numpy_scalar_types.add(dtypes.uint2)

if dtypes.int1 is not None:
  assert dtypes.uint1 is not None
  numpy_scalar_types.add(dtypes.int1)
  numpy_scalar_types.add(dtypes.uint1)

core.literalable_scalar_types.update(numpy_scalar_types)

array_types: set[type] = {literals.TypedNdArray, np.ndarray} | numpy_scalar_types


def masked_array_error(*args, **kwargs):
  raise ValueError(
      "numpy masked arrays are not supported as direct inputs to JAX functions."
      " Use arr.filled() to convert the value to a standard numpy array.")

core.pytype_aval_mappings[np.ma.MaskedArray] = masked_array_error


def _make_shaped_array_for_numpy_array(x: np.ndarray) -> ShapedArray:
  dtype = x.dtype
  dtypes.check_valid_dtype(dtype)
  return ShapedArray(x.shape, dtypes.canonicalize_dtype(dtype), sharding=None)

core.pytype_aval_mappings[np.ndarray] = _make_shaped_array_for_numpy_array
core.pytype_aval_mappings[literals.TypedNdArray] = lambda x: x.aval


def _make_shaped_array_for_numpy_scalar(x: np.generic) -> ShapedArray:
  dtype = np.dtype(x)
  dtypes.check_valid_dtype(dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix to_iterable_with_keys to return ([(key, child), ...], aux_data) where the first element is a list/tuple of 2-tuples
2. If using register_pytree_node, ensure its to_iterable returns an iterable of children and not e.g. a count or None
3. Print repr of what your hook currently returns and compare with jax.tree_util.register_pytree_node docs
4. Add a unit test that calls jax.tree_util.tree_flatten_with_path(instance) for every registered custom node

Example fix

# before
def _iter(obj):
    return (None, obj.__dict__)  # first element not iterable

# after
def _iter(obj):
    return ([(k, v) for k, v in obj.__dict__.items()], None)

import jax.tree_util as jtu

def check_to_iterable(node):
    out = node_to_iterable_with_keys(node)  # your registered hook
    assert isinstance(out, tuple) and len(out) == 2
    key_leaf_pairs, aux = out
    try:
        iter(key_leaf_pairs)
    except TypeError:
        raise ValueError('key_leaf_pairs must be iterable')
    for pair in key_leaf_pairs:
        assert isinstance(pair, tuple) and len(pair) == 2

def has_valid_flatten_hook(obj) -> bool:
    try:
        jtu.tree_flatten_with_path(obj)
        return True
    except (ValueError, TypeError):
        return False

When it happens

Trigger: Registering a custom PyTree node via jax.tree_util.register_pytree_node (which internally registers to_iterable_with_keys) whose to_iterable/to_iterable_with_keys function returns something like (None, metadata) or (count, metadata) instead of an iterable of (key, leaf) pairs; error surfaces on the first tree.flatten/tree.map over an instance.

Common situations: Porting a PyTorch/other-framework container class to JAX, porting old setattr/getattr-based flatten logic, or upgrading JAX versions where the key-returning to_iterable_with_keys protocol became mandatory and the user's hook still returns the older (children, metadata) shape with children replaced by a non-iterable.

Related errors

• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32
• The names should be exclusive and should not intersect in `n
• The return value of the policies should be a boolean. Got: {
• prevent_cse must be a bool or tuple of bools, got {type(prev

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