jax-ml/jax

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

Code / MessageTypeSeverityTags
Concatenating arrays with strided layout is only supported…
exception error jax, mosaic-gpu, not-implemented, wgmma, layout-limitation
Delta orthogonal initializer requires a 3D, 4D or 5D shape.
exception error jax, initializer, shape-validation
dimensions are not unique
validation error jax, squeeze, duplicate-dims, validation
Encountered dimension variable
exception error jax, export, shape-polymorphism, dimension-variables
inline_mgpu_p return type tree mismatch
validation error jax, pallas, pytree, return-type-mismatch
input type mismatch for
validation error jax, type-mismatch, tracing, custom-primitive
doesn't support weights. Only method 'inverted_cdf'…
validation error jax, numpy, quantile, weights, not-implemented
Only byte-aligned shapes are supported. Got shape
exception error jax, pallas, mosaic-gpu, shared-memory, alignment, shape-validation
Preemption sync manager should only be initialized once.
exception error jax, distributed, preemption, singleton
Query, key and value should have same sharding.
validation error jax, sharding, gspmd, multi-device, cudnn
should have the same core axis name and number of cores as…
validation error jax, pallas, tpu, sparsecore, mesh-compat
sharding with memory_kind is not allowed. Please use…
exception error jax, sharding, memory-kind, device-put
size and fill_value arguments cannot be used in three-term…
exception error jax, where, invalid-argument, api-misuse
Swizzled dims cannot be sliced
validation error jax, pallas, slicing, swizzle
the gradient for , which is Ref-typed, can't be returned as…
validation error jax, vjp, refs, autodiff
The input arguments to the custom_vmap-decorated function
validation error jax, custom-vmap, kwargs, argument-binding
The is_deleted() method was called on
exception error jax, concretization, tracer, memory-management, is-deleted
{_vjp_too_many_args(name, len(extra_args) + 1)}
validation error jax, vjp, arity, argument-count
All input arrays must have the same shape. Got
validation error jax, stack, shape-mismatch
block_q must be a multiple of
exception error jax, pallas, tpu, splash-attention, segment-ids, block-size
CSC must have ndim=2; got
exception error jax, sparse, csc, shape-validation
Got invalid memory kind
validation error jax, sharding, memory-kind
Index along axis is out of bounds for shape
validation error jax, mosaic-gpu, memref, index-out-of-bounds, indexerror
index is out of bounds for axis
validation error jax, empty-array, out-of-bounds, indexing
index_map returned a value of type
validation error jax, pallas, index-map, scalar
interior padding in padding_config must be nonnegative, got…
validation error jax, pad, negative-padding, validation
Invalid for
exception error jax, sparse, bcoo, shape-validation, nse
Non-divisible unfold:
validation error jax, mosaic-gpu, memref, reshape, non-divisible-shape
Only TiledLayout supports reductions.
validation error jax, pallas, mosaic-gpu, layout, reduction, unsupported-type
takes a scalar pred as argument, got
exception error jax, checkify, type-error, scalar-predicate
ragged_dot_general requires the group count (last dimension…
validation error jax, ragged-dot-general, dynamic-shapes, jit
shard_map requires a non-empty mesh. Got
exception error jax, shard-map, mesh, empty-mesh, distributed
Strided slices unsupported. Got stride
exception error jax, pallas, slicing, strides, tiling
unexpected JAX type (e.g. shape/dtype) for argument to VJP…
validation error jax, vjp, dtype, shape-mismatch
Unrecognized index type
validation error jax, internal, invariant
When changing to a larger dtype, its size must be a divisor…
exception error jax, view, dtype, alignment
Cannot `wait_send` on a local copy.
exception error jax, pallas, mosaic, dma, remote-copy, wait-send
custom_call backend_config unexpected type
exception error jax, mlir, custom-call, backend-config, type-error
Custom JVP rule for function must produce primal and…
exception error jax, custom-jvp, pytree-mismatch
exp1 does not support complex-valued inputs.
exception error jax, exp1, complex-numbers, unsupported-dtype, exponential-integral
Incompatible FragmentedArray shapes
exception error jax, mosaic-gpu, shape-mismatch, registers, pointwise
interpret with dynamic grid bounds unsupported
exception error jax, pallas, jvp, autodiff, dynamic-grid, notimplementederror
Invalid dtype for `addupdate`. Ref dtype
validation error jax, dtype-mismatch, state-primitives, in-place-update
method argument to `poisson` must be one of
validation error jax, random, poisson, method, input-validation
nan
exception error jax, nan, numerical-debugging
Non-divisible tiling:
validation error mosaic-gpu, tiling, divisibility
padding_value must be a scalar; got
validation error jax, pad, scalar-required, shape-validation
PRNG with name already registered
exception error jax, prng, registry, duplicate-registration
Python int too large to convert to int32
validation error pytree, namedtuple, jax, tree-flatten
Shape too large: >
validation error jax, pallas, tpu, random, philox, shape-limit
The length of `a` along the last axis must be at least 2…
exception error jax, companion-matrix, polynomial, input-validation
Unsupported block dim type
validation error jax, pallas, blockspec, type-error
You can't use two different TensorCoreMeshes.
validation error jax, pallas, tpu, mesh, tensorcore, multi-core
must be a multiple of .
validation error jax, pallas, tpu, splash-attention, block-size
Expected metadata dtype to be uint2, got
validation error jax, pallas, mosaic-gpu, sparse, tcgen05, dtype, validation
Explicitly requested dtype
console warning jax, dtype, precision, x64
Input 'T' must be square.
exception error jax, rsf2csf, schur-form, shape-validation
jax.numpy. with a symbolic number of sections is not…
exception error jax, split, symbolic-dim, tracing
Mesh context manager is disabled.
error_code error jax, mesh, deprecation, context-manager, config-flag
mode must be 'largest' or 'smallest', got
validation error jax, top-k, invalid-argument, value-error, enum-validation
No JVP defined for custom_jvp function
exception error jax, custom-jvp, autodiff, missing-rule
Only supported for element type , but got TMEM packing of
exception error mosaic, gpu, tcgen05, packing, not-implemented, jax
out argument of .reduce()
validation error jax, ufunc, reduce, immutable-arrays
Out-of-bounds block index
exception error pallas, mosaic-gpu, interpret-mode, out-of-bounds, input-block
Passing complex start/stop/step to jnp.arange is no longer…
validation error jax, arange, complex, breaking-change, v0-10
`preferred_element_type` must have the same signedness as…
exception error jax, dot-general, signedness, dtype-validation
primal and tangent arguments to jax.jvp do not match…
validation error jax, jvp, dtype, autodiff, float0
Run out of characters for batch dimension in einsum.
exception error jax, einsum, vmap, resource-limit
Scatter only supports VectorSubcoreMesh, got
exception error jax, pallas, sparsecore, mesh, scatter
The argument `side` must be either 'right' or 'left'.
exception error jax, argument-validation, polar-decomposition
Unsupported input type to jax.numpy.i0
validation error jax, i0, bessel, dtype, complex
Advanced indexers are not supported on TPU
validation error jax, pallas, tpu, interpreter, advanced-indexing
Arguments to jax.scipy.linalg.block_diag must have at most…
validation error jax, linalg, block-diagonal, shape-mismatch
Arguments to rng_uniform must have identical dtypes, got
validation error jax, rng-uniform, dtype-mismatch
axes out of range for array with
exception error jax, nonzero, axes, argument-validation
`buffer_callback` not supported on
exception error jax, backend, tpu, unsupported-platform, buffer-callback
dtype argument to `chisquare` must be a float dtype, got
validation error jax, random, chisquare, dtype, input-validation
dtype argument to `gamma` must be a float dtype, got
validation error jax, random, gamma, dtype, input-validation
Expected indexer to have exactly
validation error jax, pallas, indexing, rank-mismatch
Expected tiling to be at most rank of shape. Got tiling
validation error mosaic-gpu, tiling, shape-validation
get_global only supports semaphores, got
exception error jax, pallas, get-global, semaphore
hyp1f1 does not support complex-valued inputs.
exception error jax, scipy, hyp1f1, complex-dtype, unsupported-operation
linear_transpose only supports [float or complex] -> [float…
validation error jax, linear-transpose, dtype
megacore_mode must be one of ['kv_head', 'batch', None]
validation error jax, pallas, tpu, paged-attention, megacore, invalid-argument
Output layout must match the accumulator layout
exception error jax, mosaic, gpu, reduction, layout
predicate not supported
exception error pallas, mosaic-gpu, interpret-mode, predicate, not-implemented
`preferred_element_type` must not be narrower than the…
validation error jax, dot-general, preferred-element-type, dtype-narrowing
Range is entirely out of bounds for shape .
validation error jax, pallas, mosaic, index-error, out-of-bounds, interpret-mode
reductions require axes to be (0,) on SparseCore, but got
exception error jax, pallas, sparsecore, reduction, axes
Reductions with constants not supported.
exception error jax, pallas, triton, reduction, closures
The iteration bound corresponding to the core-parallel…
exception error tpu, pallas, dynamic-shape, static-bounds, jax
Transforms not supported for matmul_pop.
exception error jax, pallas, tpu, matmul, not-implemented
Unknown resize method
validation error jax, image-resize, invalid-argument
Unsupported packing: for element type
exception error mosaic, gpu, tcgen05, packing, not-implemented, jax
Unsupported TMA index shape
validation error jax, mosaic-gpu, tma, async-load, indices, gather
Unsupported WGMMA features with A in registers
validation error jax, mosaic-gpu, wgmma, api-misuse, registers
all_to_all must be used within a mapped context like vmap…
exception error jax, all-to-all, mapped-context
Argument `n` to polygamma must be of integer type. Got dtype
validation error jax, scipy-special, polygamma, dtype-validation
Cannot partition grid over dynamic number of cores.
exception error jax, pallas, tpu, grid-partition, static-value
conv_general_dilated feature_group_count must divide lhs…
validation error jax, lax, convolution, feature-group-count