{"record":{"id":"32b3a4e616cf84a1","repo":"jax-ml/jax","slug":"invalid-out-shape-type-type-out-shape","errorCode":null,"errorMessage":"Invalid out_shape type: {type(out_shape)}","messagePattern":"Invalid out_shape type: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":1706,"sourceCode":"            shape=out_shape.shape, dtype=out_shape.dtype,\n            sharding=jax_core.get_cur_mesh_sharding(),\n            manual_axis_type=out_shape.manual_axis_type)\n      return jax_core.ShapedArray(\n          shape=out_shape.shape, dtype=out_shape.dtype,\n          sharding=jax_core.get_cur_mesh_sharding())\n    case jax_core.ShapedArray():\n      return out_shape\n    case MemoryRef():\n      return out_shape.get_array_aval()\n    case hijax.HiType():\n      return out_shape\n    case _:\n      if type(out_shape) in _out_shape_to_aval_mapping:\n        return _out_shape_to_aval_mapping[type(out_shape)](\n            out_shape\n        )\n      if not (hasattr(out_shape, \"shape\") and hasattr(out_shape, \"dtype\")):\n        raise ValueError(f\"Invalid out_shape type: {type(out_shape)}\")\n      return jax_core.ShapedArray(shape=out_shape.shape, dtype=out_shape.dtype)\n","sourceCodeStart":1688,"sourceCodeEnd":1708,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L1688-L1708","documentation":"While converting a declared out_shape into a JAX abstract value (aval), jax found an object that is neither a ShapeDtypeStruct, a known mapped type, nor a duck-typed object exposing both .shape and .dtype. _convert_out_shape_to_aval therefore rejects it with ValueError. The API requires out_shapes to look like shape/dtype descriptors.","triggerScenarios":"Passing an int, tuple, numpy dtype, string, or arbitrary object as out_shape/out_shapes entry to a pallas/shard_map-style API that routes through _convert_out_shape_to_aval (jax/_src/pallas/core.py:1706).","commonSituations":"Passing a bare shape tuple (8, 8) instead of ShapeDtypeStruct((8,8), dtype); passing a custom container that lacks .shape/.dtype attributes; refactoring where out_shape became a dataclass without those attribute names.","solutions":["Wrap the shape in jax.ShapeDtypeStruct(shape, dtype) before passing it as out_shape","If using a custom class, give it .shape and .dtype properties so duck-typing succeeds","Pass jax.core.ShapedArray directly, which is in the known mapping","Check you are not accidentally passing dtype or axis names where out_shape is expected"],"exampleFix":"# before\nkernel = pallas_kernel(fn, out_shape=(8, 8))  # tuple has no .shape/.dtype\n\n# after\nout = jax.ShapeDtypeStruct((8, 8), jnp.float32)\nkernel = pallas_kernel(fn, out_shape=out)","handlingStrategy":"type-guard","validationCode":"def validate_out_shape(o):\n    assert hasattr(o, 'shape') and hasattr(o, 'dtype'), f'out_shape {type(o)} must expose .shape and .dtype'","typeGuard":"from typing import Any\ndef is_valid_out_shape(o: Any) -> bool:\n    return hasattr(o, 'shape') and hasattr(o, 'dtype')","tryCatchPattern":"try:\n    aval = convert_out_shape(o)\nexcept ValueError as e:\n    if 'Invalid out_shape type' in str(e):\n        o = jax.ShapeDtypeStruct(o.shape, o.dtype)  # normalize then retry\n        aval = convert_out_shape(o)\n    else:\n        raise","preventionTips":["Standardize on jax.ShapeDtypeStruct for every declared out_shape","Validate out_shape lists with a small helper before passing to pallas/shard_map APIs","Add type annotations (ShapeDtypeStruct) so mypy catches bad types early"],"tags":["jax","pallas","out-shape","shapedtypestruct","input-validation"],"backgroundTag":"invalid-output-shape-declaration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}