ErrLookup › jax-ml/jax

jax-ml/jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more · Python · 870 source files

Analyzed at 620b2c5563 on 2026-09-02. 2226 documented errors.

Code / MessageTypeSeverityTags
No constant handler for type
exception error jax, constants, type-error, tracing
JAX encountered invalid PRNG key data: expected key_data to…
exception error jax, prng, typeerror, input-validation
Buffer callbacks do not support transpose. Please use…
exception error jax, autodiff, gradient, callback, buffer-callback
Expected tuple, got .
validation error jax, pytree, type-mismatch, tuple
cannot accept args with unreduced_kind= . Expected…
exception error jax, shard-map, spmd, collectives, internal-api
function carry input and carry output must have the same…
validation error jax, pytree, scan, while-loop, carry-structure
has no attribute
exception error jax, tracer, attributeerror, jit, duck-typing
K tile stride must be a multiple of 16
validation error gpu, mosaic, tcgen05, sparse, shared-memory, alignment
Scan carry input and output got mismatched varying manual…
validation error jax, shard-map, vma, scan, spmd, multi-device
buffer_callback callback must not return any values.
exception error jax, callback, host-callback, return-value, buffer-callback
Unsupported copy: ->
exception error jax, mosaic-gpu, gpu-kernel, shared-memory, not-implemented
Only take_along_axis-like gathers supported
exception error jax, pallas, tpu, gather, not-implemented
Denominator should be reduced along the same axes numerator…
exception error jax, sharding, lax, division, distributed
Source and destination layouts aren't compatible for a…
exception error jax, mosaic-gpu, layout, broadcast, fragmented-array
scan number of arguments doesn't match the number of jaxpr…
validation error jax, internal, scan, jaxpr, argument-count
shard_map out_specs rank error (msg from _spec_rank_error)
exception error jax, shard-map, out-specs, rank, distributed
requires all arguments to have matching type. Got key type…
validation error jax, sharding, shard-map, pvary, random
Invalid order ' ' for vector norm.
exception error jax, linalg, norm, type-validation
Data type bitcast is only supported from i8 to other types.
exception error jax, pallas, mosaic-gpu, bitcast, dtype, not-implemented
shard_map out_specs vma error (msg from _inout_vma_error)
exception error jax, shard-map, out-specs, vma, replication
in_axes was not specified when axis_name=
exception error jax, shard-map, smap, axis-types, in-specs
pvary is a invariant->variant collective. This means that…
validation error jax, collectives, mesh, spmd, pvary
pallas_call requires all mesh axes to be Manual, got
exception error jax, pallas, pallas-call, sharding, mesh, manual-axes
x and n must be of integer type; got x.dtype=
exception error jax, scipy, multinomial, dtype, statistics
vmapping pallas_call with no arguments.
exception error jax, pallas, vmap, batching, not-implemented
in_specs passed to shard_map
exception error jax, shard-map, sharding, partition-spec, distributed
eigh_tridiagonal(..., select='v') is not implemented
validation error jax, eigenvalues, not-implemented, linalg
Mask function must return a boolean-valued array, but got
validation error jax, pallas, tpu, splash-attention, dtype, mask
Argument ' ' of type is not a valid JAX type.
validation error jax, type-validation, argument-validation
function carry input and carry output must have equal…
validation error jax, dtype-mismatch, scan, while-loop, shard-map, vma
method='qdwh' only supports mxn matrices where m < n where…
exception error jax, polar-decomposition, qdwh, shape-mismatch
does not have an attribute
exception error mosaic-gpu, mlir, missing-attribute, layout-inference
transpose_solve required for backwards mode automatic…
validation error jax, autodiff, custom-linear-solve, transpose
Formatting arguments to checkify.check need to be PyTrees…
exception error jax, checkify, typeerror, assertions, jit
'order' argument to argsort is not supported.
validation error jax, argsort, numpy-compat, unsupported-argument, sorting
Parameter is not a Jaxpr or sequence of Jaxprs
validation error jax, pallas, jaxpr, control-flow, interpreter
Value of type is not compatible with the Array API.
exception error jax, tracer, numpy-interop, array-api, nep-35
copy_gmem_to_smem without a barrier is only supported on…
exception error mosaic-gpu, pallas, gpu, barrier, tma, hopper
MPMD map with more than one mesh requires scratch_type to…
exception error jax, pallas, mpmd, tpu, scratch-memory, multi-mesh
check_error takes an Error as argument, got type
exception error jax, checkify, typeerror, api-misuse
packed cannot be specified if layout is specified.
validation error jax, pallas, mosaic-gpu, tmem, layout, api-misuse
Stacking only supported when the block size along the stack…
exception error jax, pallas, fusion, block-spec, shape-mismatch
dirichlet requires alpha.ndim >= 1, got alpha.ndim ==
validation error jax, random, dirichlet, input-validation, shape
open an issue at https://github.com/google/jax !!
error_code error jax, autodiff, custom-linear-solve, not-implemented
pure_callback only supports SingleDeviceSharding, but got
exception error jax, sharding, pure-callback, single-device-sharding
strides=
validation error jax, mosaic-gpu, memref, strides, layout
only accepts inputs that are unreduced. Got
exception error jax, sharding, collectives, validation, mesh
Nonsymmetric eigendecomposition requires cusolver 11.7.1 or…
error_code error jax, cuda, cusolver, eig, version-mismatch, gpu
Out-of-bounds read of
validation error jax, pallas, tpu, out-of-bounds, interpret-mode
jax.pure_callback failed to find a local CPU device to…
exception error jax, environment, pure-callback, cpu-device, jax-platforms
Ordered IO effects not supported in vmap.
error_code error jax, vmap, ordered-io, debug-print, while-loop, batching
Subclasses should override this method
exception error mosaic-gpu, not-implemented, abstract-method, memref-transform, subclassing
The Pallas TPU lowering currently requires that the last…
validation error jax, pallas, tpu, alignment, dynamic-shapes
Not all transforms could be handled. Remaining transforms
exception error jax, pallas, mosaic-gpu, transforms, not-implemented
The collective size ( ) must divide the slice shape along…
validation error jax, mosaic-gpu, async-copy, cluster-partitioning, shape-validation
does not support and
exception error jax, pallas, mosaic-gpu, dtype, binary-op, not-implemented
Unordered IO effects not supported in while_loop with…
error_code error jax, vmap, while-loop, io-effects, control-flow
associative scan over axis of non-constant size
error_code error jax, associative-scan, dynamic-shape, shape-polymorphism
`sharding` passed to `convert_element_type` can only…
exception error jax, sharding, dtype-conversion, lax
Unsupported dtype
validation error jax, linalg, eig, dtype, gpu
`JaxprInputEffect` does not have a corresponding jaxpr…
validation error jax, jaxpr, effects, internal-invariant, mutable-arrays
The Pallas TPU lowering currently supports in memory space…
validation error jax, pallas, tpu, block-mapping, memory-space
Invalid dtype for `swap`. Ref dtype
validation error jax, dtype-mismatch, state-primitives, swap
pages_per_compute_block must be divisible by pages per…
validation error jax, pallas, tpu, paged-attention, layout-validation
Cannot interpret ' ' as a data type.\n\nDid you…
exception error jax, shape-validation, dtype, typo-detection
Failed to find assignment for logical_axis_index
exception error jax, mesh, tpu, topology, not-implemented
in op must be sorted; got
exception error jax, lax, gather, scatter, dimension-numbers
Reduced axes can only refer to mesh axes that is of type…
validation error jax, sharding, mesh, distributed, named-sharding
callbacks are only supported in spmd computations when all…
exception error jax, spmd, sharding, mesh, pure-callback, manual-axes
Existing output dtype must match preferred_element_type.
validation error jax, pallas, tpu, megablox, dtype-mismatch
n must be a non-negative integer.
exception error jax, scipy, bernoulli, static-argument, input-validation
Please file an issue at https://github.com/jax-ml/jax/issues
exception error jax, vmap, collectives, not-implemented, all-gather
Size mismatch for group
validation error jax, einshape, einsum-notation, shape-mismatch, grouped-reshape
All WGStridedFragLayout arrays must have the same vec_size…
validation error mosaic-gpu, fragmented-array, layout-mismatch, concat, validation
The ' ' method is not available on .
exception error jax, tracer, block-until-ready, async, attribute-access
Gather index leaf dimension must be within
exception error jax, lax, gather, index-vector-dim
CollapseShapeOp with non-contiguous strides is not…
validation error jax, mosaic-gpu, memref, collapse-shape, strides, layout-inference
Expected kind to be a dtype, string, or tuple; got
validation error jax, dtype, isdtype, argument-validation
should divide .
validation error jax, tpu, splash-attention, shape-mismatch, pallas
The input arguments to the custom_jvp-decorated function
validation error jax, custom-jvp, kwargs, argument-binding
Value of type is not convertible to float.
exception error jax, tracer, type-conversion, jit, float
vmap spmd_axis_name cannot appear in shard_map in_specs
validation error jax, vmap, shard-map, spmd, partition-spec
Can't print the type
exception error gpu, mosaic, debug-print, unsupported-type, dtype
Array.__contains__: unsupported operand type
exception error jax, membership, typeerror, contains
Unknown attribute
exception error mosaic-gpu, invalid-argument, attribute-dispatch, internal-api
Cannot permute last two dimensions with leading dimensions.
exception error jax, pallas, transpose, permutation, not-implemented
Invalid value received for the sharding argument. Expected…
validation error jax, sharding, type-validation
The 'sharding' attribute is not available on
exception error jax, tracer, sharding, attribute-access, distributed
Not implemented: can not be XLA fully tiled.
validation error jax, pallas, tpu, attention, head-tiling, not-implemented
has of an unexpected type
exception error mosaic-gpu, mlir, type-mismatch, attribute-type, layout-inference
multi-platform lowering for buffer_callback
exception error jax, xla, multi-platform, lowering, export
Pointer address space
validation error gpu, mosaic, address-space, pointer, memref, llvm
The SMEM tiles must be contiguous
validation error gpu, mosaic, tcgen05, strides, shared-memory, contiguity
Element with is not supported.
validation error jax, pallas, mosaic, pipeline, blockspec, padding
Invalid set in op; valid range is
exception error jax, lax, gather, scatter, dimension-out-of-range
JAX encountered invalid PRNG key data: expected…
exception error jax, prng, ndim, typeerror
unreduced/reduced can only be passed to
exception error jax, shard-map, manual-mode, unreduced, experimental
Effects not supported in `while
error_code error jax, effects, while-loop, debug-print, not-implemented
Only unit strides are supported but got
validation error jax, mosaic-gpu, memref, layout-inference, strides, not-implemented
Sharded hypothesis runner does not support `data()` inside…
exception error jax, testing, hypothesis, sharding