jax-ml/jax

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

Code / MessageTypeSeverityTags
operand to reduce_window_sum must have a number dtype, got
validation error jax, dtype, reduce-window, type-validation
Predicate not supported for no-complete arrive
validation error jax, mosaic-gpu, barrier, predicate, not-implemented
Tile shape must only have positive sizes, got
validation error mosaic, fragmented-array, tiling, validation
Unknown FFT type value
validation error jax, fft, type-error
unreduced for min is not yet supported.
exception error jax, sharding, notimplementederror, distributed
Unsupported aval type
exception error jax, checkify, shard-map, custom-aval, extensibility
A scale shape mismatch: expected
validation error gpu, mosaic, tcgen05, shape-mismatch, block-scaling
cannot handle multidimensional aweights
exception error jax, covariance, weights, shape-validation
Character '.' must be used inside ellipsis '...'
validation error jax, sharding, ellipsis, rule-syntax
covariance does not have dimension
exception error jax, scipy, kde, covariance, shape-validation
dot_general requires rhs contracting dimensions to be…
exception error jax, dot-general, dimension-numbers, contraction
fill_value argument to indexed get() must be a scalar
exception error jax, indexing, fill-value, scalar
Found inconsistent memory spaces in multiref
exception error jax, memory-space, multiref
index_dtype must be an integer type, but got
validation error jax, argmin, dtype, type-error
inline_mgpu in a single-warp context only supports scalar…
validation error jax, pallas, warp-semantics, scalar-return
lax.scan: f argument should be a callable.
validation error jax, scan, typeerror, callable, api-misuse
M must be a multiple of 16 and <= 256, got
validation error tcgen05, matmul, mosaic, gpu, shape-validation
ragged_dot vmap over any dim but 0 - NYI
exception error jax, ragged-dot, vmap, batching, not-implemented
structure of output value and output batching specification…
validation error jax, custom-vmap, pytree, vmap
vmem_limit_bytes must be an int: provided with a
exception error tpu, pallas, vmem, type-error, config
A sparse metadata dtype mismatch: expected i2, got
validation error gpu, mosaic, tcgen05, sparsity, dtype, metadata
Arguments to rng_uniform must be scalars; got shapes
validation error jax, rng-uniform, scalar-requirement, shape-validation
Array layout mismatch: expected
validation error jax, pallas, inline-mgpu, layout-mismatch
bfloat16 top_k is not supported on TPUv5 or older
exception error jax, pallas, tpu, top-k, hardware-version, bfloat16
Either none or both scales should be provided
validation error tcgen05, matmul, mx-format, argument-validation
Expected TiledLayout, got
exception error jax, mosaic, gpu, smem, layout
Input/output tiling mismatch when attempting to collapse a…
validation error jax, mosaic-gpu, collapse-shape, tiling, consistency
Only slicing with static indices allowed.
exception error jax, mosaic, gpu, vector, static-index
reduce_window batching is not implemented for initial values
exception error jax, vmap, batching, not-implemented, reduce-window
scan got `length` argument of
validation error jax, scan, length-mismatch, shape-validation
Scratch operands are not supported by the Mosaic GPU…
exception error pallas, mosaic-gpu, mpmd-map, scratch, not-implemented
The emit_pipeline body function must return None.
exception error jax, pallas, kernel-body, validation
The 'out' argument to jnp.nanstd is not supported.
validation error jax, numpy, nanstd, out-argument
todense_transpose for
exception error jax, sparse, autodiff, todense, csr
Unsupported index dtype
validation error jax, cuda, cusparse, sparse, dtype
Unsupported mode
validation error jax, nn, float8, config, enum-argument
Unsupported operand type
exception error jax, mosaic, mma, unsupported-dtype
Arguments to batch_matmul must have same ndim, got
validation error jax, lax, batch-matmul, rank-mismatch, value-error
Barrier does not support slice with `stride != 1`
validation error mosaic-gpu, pallas, barrier, indexing, slice
distributed.initialize should only be called once.
exception error jax, distributed, double-init, lifecycle
dot_general requires rhs batch dimensions to be distinct…
exception error jax, dot-general, dimension-numbers, batch-dims
dtype argument to `logistic` must be a float dtype, got
validation error jax, random, logistic, dtype, input-validation
Expanding tiled dimensions is not supported.
exception error jax, mosaic-gpu, expand-shape, tiling, not-implemented
Factor in is not used in the rule
validation error jax, sharding, unknown-factor, validation
factor_sizes must be a dict of str to int, but got
validation error jax, sharding, type-error, factor-sizes
for transpose support, subclass
validation error jax, autodiff, transpose, not-implemented
Incompatible shapes for Sylvester equation: A
exception error jax, scipy, linalg, shape-mismatch, sylvester
Input layouts do not match yield layouts
validation error mosaic, gpu, scf, while-loop, layout-mismatch
Invalid shape for other
exception error jax, splash-attention, mask, shape-mismatch
is_signed must be specified for integer types
exception error jax, mosaic, mma, dtype, signedness
No active profiler server.
validation error jax, profiler, lifecycle
`pallas` reduce operations only support one reduce axis.
exception error jax, pallas, triton, reduction, axes
Per-operand conv precision unsupported
exception error jax, pallas, tpu, convolution, precision
rule must be a str, but got
validation error jax, sharding, type-error, rule-syntax
{str(exc)}
exception error jax, checkify, assertion, debugging
Swap does not support storing to
exception error jax, pallas, tpu, sparsecore, memory-space, swap
There is no -> in rule
validation error jax, sharding, missing-separator, rule-syntax
trace_value requires i32 or f32, got
exception error jax, pallas, tpu, debugging, dtype
Unsupported transforms for ACC
exception error jax, pallas, tcgen05, transform, not-implemented
Weights shape must match reduction axes
validation error jax, numpy, quantile, weights, shape-mismatch
Barrier offset is out of bounds
validation error jax, mosaic-gpu, barrier, index-out-of-bounds
debug_print only supports printing of scalar values, or a…
exception error jax, pallas, debug-print, arrays
Dynamic grid bounds not (yet) supported in GPU interpret…
exception error pallas, mosaic-gpu, interpret-mode, dynamic-grid, not-implemented
for select='i', select_range must be specified.
validation error jax, linalg, tridiagonal, missing-argument
gaussian_kde does not support complex coordinates
exception error jax, scipy, kde, complex-dtype, not-implemented
gpu.launch op not found.
validation error mosaic, gpu, gpu-launch, pass-pipeline, mlir
Integers cannot be raised to negative powers, got…
exception error jax, lax, integer-pow, negative-exponent, type-error
Invalid shape for kv segment_ids
validation error jax, pallas, tpu, splash-attention, segment-ids, shape
length of padding_config must equal the number of axes of…
validation error jax, pad, config-mismatch, rank-error
memref.LoadOp does not support transforms
exception error jax, mosaic-gpu, memref-load, transforms, not-implemented
mode must be one of ['same', 'full', 'valid']
exception error jax, scipy, signal, argument-validation
new int config value must be None or of type int, got
validation error jax, config, valueerror, int
Only float32 and float64 inputs are supported as inputs to…
validation error jax, linalg, tridiagonal, dtype
orthogonal initializer requires at least a 2D shape
exception error jax, nn, initializers, shape-validation
program id was requested but no grid was provided.
exception error jax, pallas, mosaic, grid, program-id, lowering
scan got values with different leading axis sizes
validation error jax, scan, shape-mismatch, length-inference
Size mismatch for in/out transforms. In transforms
validation error jax, mosaic-gpu, transpose, transforms, count-mismatch
Sparse MMA unsupported for f32
exception error jax, mosaic, gpu, sparse, dtype, mma, tcgen05
sparse rule for with non-positive exponent is not…
exception error jax, sparse, sparsify, power, densifying
static_slice requires mode='promise_in_bounds' or…
validation error jax, indexing, mode, api-misuse
Sum of sizes must be equal to dimension of the operand shape
validation error jax, split, size-mismatch, shape-validation
The 'out' argument to jnp.take is not supported.
validation error jax, numpy-compat, out-arg, immutable
Trend type must be 'linear' or 'constant'.
exception error jax, scipy, detrend, invalid-argument-value
Unsupported A register array layout
validation error jax, mosaic-gpu, wgmma, layout, fragmented-array
varying and reduced cannot have common mesh axes. Got…
validation error jax, sharding, manual-axis-type, validation
Async load only supports TMEM refs
validation error jax, pallas, tmem, memory-space, blackwell
Attempting to deallocate allocation with key
exception error jax, pallas, mosaic, shared-memory, key-collision, interpret-mode
axis argument out of range
validation error jax, top-k, axis, out-of-bounds
c_api argument to register_xla_transform_c_api is not a…
validation error jax, xla, pjrt, pycapsule
Cannot concatenate an empty list of vectors
validation error mosaic-gpu, vector-concat, empty-input
dynamic_slice requires mode='promise_in_bounds' or…
validation error jax, dynamic-slice, mode, api-misuse
Expected metadata to be 3-dimensional (M, K // 4, 2), but…
validation error jax, pallas, mosaic-gpu, sparse, tcgen05, shape, rank
Got ; expected 'auto', 'fft', or 'direct'.
exception error jax, scipy, convolution, invalid-argument-value
jax.numpy.quantile does not support overwrite_input=True or…
validation error jax, numpy, quantile, out-argument, overwrite-input
k must be a scalar or a rank-1 array of length 1 or m.
validation error jax, polyint, shape-mismatch
MMA rhs tiling does not fit swizzle
validation error jax, pallas, tcgen05, swizzle, tiling, layout
Ordered debug_print is not supported on Pallas.
exception error jax, pallas, debug-print, ordered
reduce_window got inconsistent window_dilation and…
validation error jax, shape-validation, reduce-window, dilation
reduced_vary_cast is a Reduced->Varying collective. This…
exception error jax, axis-name, named-axes
tcgen05_mma only allows arriving on a Barrier that orders…
exception error pallas, mosaic-gpu, tcgen05, mma, barrier, tensor-core