jax-ml/jax

Documented errors, page 12 of 23. Back to jax-ml/jax

Code / MessageTypeSeverityTags
unsupported eviction policy
exception error jax, triton, pallas, load, eviction-policy, validation
Axis is out of bounds for grid
validation error jax, pallas, grid, index-out-of-bounds
Axis name not found in grid.
exception error jax, pallas, triton, grid, axis-index
dot_general requires lhs batch dimensions and rhs batch…
validation error jax, dot-general, shape-mismatch, batch-dims
expected a 1-d array for weights
validation error jax, numpy, polynomial, polyfit, weights, shape-validation
invalid flag value
exception error tpu, pallas, custom-call, config-serialization, jax
jnp.linalg.cond: for
exception error jax, numpy, linalg, condition-number, shape-validation
must provide `length` to `scan`, since the leading-axis…
validation error jax, scan, hijax, length, experimental
n must be a positive power of 2; got
exception error jax, scipy, linalg, matrix, argument-validation
out argument of .accumulate()
exception error jax, ufunc, accumulate, immutable-arrays
partitions cannot overlap with reduced axes passed to…
validation error jax, sharding, partition-spec, mesh
Profile has already been started. Only one profile may be…
error_code error jax, profiler, concurrency, singleton
smap in_axes must be an int, None, jax.sharding.Infer, or a…
exception error jax, smap, shard-map, in-axes, distributed
stride_axis is out of range
exception error jax, pallas, mosaic, roll, axis-validation, shape-validation
Swap only supports scalars in SMEM.
exception error jax, pallas, tpu, sparsecore, masked-store, swap
The devices() method was called on
exception error jax, concretization, tracer, devices, placement
.transpose should return None or a dict of backward-pass…
validation error jax, autodiff, transpose, api-contract
+
exception error jax, pallas, mosaic, dtype, not-implemented, interpret-mode
All objects to concatenate must be arrays, got
validation error jax, concatenate, type-validation, dtype
cannot specify both devices and num_cores
validation error jax, pallas, tpu, mesh, mutually-exclusive-args
`dataset` input should have multiple elements.
exception error jax, scipy, kde, input-validation
in1 and in2 should have the same dimensionality
exception error jax, scipy, signal, fft, shape-mismatch
Invalid dimension range passed to collapse
validation error jax, lax, collapse, dimension-range, value-error
layouts passed to `with_layout_constraint` must be of type…
validation error jax, layout, type-validation
Lowering multiple mesh/function pairs is not supported by…
exception error pallas, mosaic-gpu, mpmd-map, not-implemented
Output layout does not match the layout of the returned…
validation error jax, pallas, layout-mismatch, mosaic-gpu
precision matrix must match data dims
exception error jax, scipy, kde, precision-matrix, shape-validation
has a memory space that is not TMEM.
validation error mosaic, gpu, tmem, memory-space, mlir
should have a different core axis name from the…
validation error jax, pallas, tpu, sparsecore, axis-naming
shard_map in_specs divisibility error (msg from…
exception error shard-map, divisibility, shape-mismatch, jax
Stores to TMEM are asynchronous operations and cannot be…
exception error jax, pallas, mosaic-gpu, tmem, blackwell, tcgen05
The current TPU does not have SparseCores
exception critical jax, pallas, tpu, sparsecore, hardware, environment
Unexpected type for value
validation error jax, pallas, inline-mgpu, type-validation
Unknown backend: ' ' requested, but no platforms that are…
exception error
VJP not implemented for select_and_gather (MaxPool) with…
error_code error jax, autodiff, vjp, pooling, dilation, not-implemented
Block size must be a multiple of the input size. Got block
exception error jax, pallas, tile, shape-mismatch, divisibility
conv_general_dilated feature_group_count must be a positive…
validation error jax, lax, convolution, feature-group-count
dtype argument to `gumbel` must be a float dtype, got
validation error jax, random, gumbel, dtype, input-validation
Encountered non-leading UntilingTransform or UnswizzleRef…
exception error pallas, mosaic-gpu, transforms, not-implemented
Expected num_classes to match the size of axis
validation error jax, nn, one-hot, pmap, distributed
hankel: c must be at least 1-dimensional, got a scalar.
exception error jax, hankel, input-validation, shape-error
Invalid shape for q segment_ids
validation error jax, pallas, tpu, splash-attention, segment-ids, shape
Invalid 'trans' value
validation error jax, linalg, solve-triangular, argument-validation
dtype should be , but got
validation error jax, nn, attention, dtype-validation
pmap in_axes: value not a valid prefix (dynamic message…
validation error jax, pmap, in-axes, pytree-prefix
primal and tangent arguments to jax.jvp must have the same…
validation error jax, jvp, pytree, autodiff
random bits array of size exceeding 2 ** 64
exception error jax, prng, philox, shape-validation, resource-limits
has no `dtype`.
exception error jax, attribute-error, dtype
slice_sizes must be less than or equal to operand shape…
exception error jax, sparse, bcoo, window-too-large, dynamic-slice
stride and stride_axis must be both specified or not.
exception error jax, pallas, mosaic, roll, argument-validation, kwargs-pairing
Unsupported reduction kind
exception error mosaic, gpu, tcgen05, reduction, internal, jax
val.dtype= != ref.dtype=
exception error jax, pallas, sparsecore, scatter, dtype-mismatch
Block sizes for dq kernel are not needed with a fused…
validation error jax, pallas, tpu, splash-attention, config-validation
Explicit sharding is not currently supported in the…
exception error jax, sharding, pallas, spmd, not-implemented
Mock topology must be of the form "<number-of-slices> x…
exception error
Only SMEM and TMEM refs are supported.
exception error jax, pallas, mosaic-gpu, ref-union, memory-space, not-implemented
Reductions over unsigned integers not implemented.
exception error jax, pallas, tpu, reduction, unsigned-integer, dtype
smap out_axes must be an int, None, or (nested) container…
exception error jax, smap, out-axes, validation
Sparse MMA not supported for M=64
exception error jax, mosaic, gpu, sparse, mma, tensor-core, tcgen05, blackwell
stream argument of array.to_device()
exception error jax, device-transfer, stream, not-implemented
an index can only have a single ellipsis ('...')
validation error jax, indexing, ellipsis
Collapsing the shape of a memref with non-contiguous…
exception error jax, mosaic-gpu, collapse-shape, strides, contiguity
Dimension mismatch
validation error jax, numpy-reference, convolution, shape-validation
DMA partial discharge add=True not yet implemented.
exception error jax, pallas, mosaic, dma, discharge, not-implemented, pipelining
JAX only supports number, bool, and string dtypes, got dtype
validation error jax, dtype, validation, numeric-types
mla requires cudnn version >= 9.10 and at least hopper arch.
exception error jax, cudnn, mla, version-mismatch, gpu-architecture
reduce_window expected init_values to be scalars but…
validation error jax, reduce-window, shape-validation, scalar
shape must have length equal to the number of dimensions of…
validation error jax, image-resize, shape-mismatch
Unsupported block shape type
validation error jax, pallas, block-spec, internal
all_to_all requires the size of the mapped axis axis_name…
exception error jax, all-to-all, shape, spmd
Array type must have a `shape` attribute, but got
validation error jax, pallas, block-spec, aval
Can't bitcast (of bitwidth ) to (of bitwidth )
validation error mosaic-gpu, bitcast, bitwidth
Cannot partition over cores without parallel grid dimensions
exception error jax, pallas, tpu, grid-partition, dimension-semantics
Cannot retrieve the architecture: no module found
validation error jax, mosaic-gpu, architecture, mlir-module
duplicate axes are not allowed
exception error jax, nonzero, axes, argument-validation
entr does not support complex-valued inputs.
validation error jax, scipy-special, entropy, complex-dtype
Expected kind to be one of
exception error jax, scipy, linalg, argument-validation
Invalid grid in grid_mapping: expected grid to have the…
exception error pallas, mosaic-gpu, grid, mesh, config-mismatch
jax.shard_map requires axis_names=
exception error jax, shard-map, axis-names, mesh, distributed
Only 2-argument concatenate is supported.
exception error jax, triton, pallas, concatenate, not-implemented
Only bitcast between types of the same bitwidth supported
validation error mosaic, gpu, bitcast, bitwidth
Replica lowering for Mosaic kernels not implemented.
exception error tpu, pallas, communication, replication, jax
series_order must be a Python integer.
exception error jax, scipy-special, log-ndtr, type-validation
shard_map in_specs rank error (msg from _spec_rank_error)
exception error shard-map, rank-mismatch, partition-spec, jax
'shift' and 'axis' arguments to roll must be scalars or 1D…
exception error jax, roll, shape-validation
take_along_axis indices must be of integer type, got
validation error jax, take-along-axis, dtype, integer-index
The ratio of contracting dim to its scale's dim size ( )…
exception error jax, scaled-dot, shape-validation
The symbolic constraints should be a sequence of strings…
exception error jax, shape-polymorphism, constraints, api-misuse
Unsupported packing
exception error mosaic, gpu, tcgen05, debug, packing, not-implemented, jax
A compound factor should contain at least two factors
exception error jax, shardy, sharding-rule, validation
Empty array
validation error gpu, mosaic, tma, empty-array, argument-validation
Expected 'key' and 'value' to have the same leading…
validation error jax, pallas, tpu, splash-attention, shape
for jvp support, subclass
validation error jax, autodiff, forward-mode, not-implemented
Group sizes must match first dimension of
exception error jax, pallas, moe, shape, gpu
Invalid = for
exception error jax, sparse, bcoo, indices, invariant
does not match expected shape
exception error jax, pallas, sparsecore, scatter, mask-shape
Mesh shape of the input
exception error shard-map, mesh-mismatch, named-sharding, jax
{non_mesh_axes}
exception error jax, pallas, mesh, device-id, axis-name, interpret-mode
Only arrays with 32-bit element types can be converted to…
exception error jax, pallas, tpu, squeeze, scalar, dtype
Only implemented for order='K'
exception error jax, not-implemented, numpy-compat, memory-order