{"record":{"id":"f9c0fbc391e3c757","repo":"jax-ml/jax","slug":"numpy-masked-arrays-are-not-supported-as-direct-in","errorCode":null,"errorMessage":"numpy masked arrays are not supported as direct inputs to JAX functions. Use arr.filled() to convert the value to a standard numpy array.","messagePattern":"numpy masked arrays are not supported as direct inputs to JAX functions\\. Use arr\\.filled\\(\\) to convert the value to a standard numpy array\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/abstract_arrays.py","lineNumber":54,"sourceCode":"} | {np.dtype(dt).type for dt in dtypes._float_types}\n\nif dtypes.int2 is not None:\n  assert dtypes.uint2 is not None\n  numpy_scalar_types.add(dtypes.int2)\n  numpy_scalar_types.add(dtypes.uint2)\n\nif dtypes.int1 is not None:\n  assert dtypes.uint1 is not None\n  numpy_scalar_types.add(dtypes.int1)\n  numpy_scalar_types.add(dtypes.uint1)\n\ncore.literalable_scalar_types.update(numpy_scalar_types)\n\narray_types: set[type] = {literals.TypedNdArray, np.ndarray} | numpy_scalar_types\n\n\ndef masked_array_error(*args, **kwargs):\n  raise ValueError(\n      \"numpy masked arrays are not supported as direct inputs to JAX functions.\"\n      \" Use arr.filled() to convert the value to a standard numpy array.\")\n\ncore.pytype_aval_mappings[np.ma.MaskedArray] = masked_array_error\n\n\ndef _make_shaped_array_for_numpy_array(x: np.ndarray) -> ShapedArray:\n  dtype = x.dtype\n  dtypes.check_valid_dtype(dtype)\n  return ShapedArray(x.shape, dtypes.canonicalize_dtype(dtype), sharding=None)\n\ncore.pytype_aval_mappings[np.ndarray] = _make_shaped_array_for_numpy_array\ncore.pytype_aval_mappings[literals.TypedNdArray] = lambda x: x.aval\n\n\ndef _make_shaped_array_for_numpy_scalar(x: np.generic) -> ShapedArray:\n  dtype = np.dtype(x)\n  dtypes.check_valid_dtype(dtype)","sourceCodeStart":36,"sourceCodeEnd":72,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/abstract_arrays.py#L36-L72","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","solutions":["Fix to_iterable_with_keys to return ([(key, child), ...], aux_data) where the first element is a list/tuple of 2-tuples","If using register_pytree_node, ensure its to_iterable returns an iterable of children and not e.g. a count or None","Print repr of what your hook currently returns and compare with jax.tree_util.register_pytree_node docs","Add a unit test that calls jax.tree_util.tree_flatten_with_path(instance) for every registered custom node"],"exampleFix":"# before\ndef _iter(obj):\n    return (None, obj.__dict__)  # first element not iterable\n\n# after\ndef _iter(obj):\n    return ([(k, v) for k, v in obj.__dict__.items()], None)","handlingStrategy":"validation","validationCode":"import jax.tree_util as jtu\n\ndef check_to_iterable(node):\n    out = node_to_iterable_with_keys(node)  # your registered hook\n    assert isinstance(out, tuple) and len(out) == 2\n    key_leaf_pairs, aux = out\n    try:\n        iter(key_leaf_pairs)\n    except TypeError:\n        raise ValueError('key_leaf_pairs must be iterable')\n    for pair in key_leaf_pairs:\n        assert isinstance(pair, tuple) and len(pair) == 2","typeGuard":"def has_valid_flatten_hook(obj) -> bool:\n    try:\n        jtu.tree_flatten_with_path(obj)\n        return True\n    except (ValueError, TypeError):\n        return False","tryCatchPattern":null,"preventionTips":["Smoke-test every register_pytree_node class with jax.tree_util.tree_flatten_with_path at registration time","Keep to_iterable and to_iterable_with_keys return shapes in sync","Add a CI test that flattens/unflattens round-trips each registered custom node"],"tags":["pytree","custom-node","jax","tree-flatten","api-contract"],"backgroundTag":"pytree-custom-node-protocol-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}