jax-ml/jax
Documented errors, page 6 of 23. Back to jax-ml/jax
| Code / Message | Type | Severity | Tags |
|---|---|---|---|
| Incompatible shapes for matrix multiplication: lhs= | validation | error | mosaic-gpu, pallas, wgmma, shape-mismatch, matrix-multiply |
| LHS and RHS must have the same number of batch dimensions… | exception | error | jax, scaled-dot, dimension-numbers |
| Only plgpu.Layout.WG_SPLAT is supported for scalar values. | exception | error | pallas, mosaic-gpu, layout, scalar |
| Profiler collected an odd number of trace events. This… | exception | error | profiling, mosaic, gpu, memory-corruption |
| __rpow__ modulo not implemented | exception | error | jax, shape-polymorphism, rpow, not-implemented |
| The is_fully_replicated property was called on | exception | error | jax, concretization, tracer, sharding, replication |
| tuple of axes | exception | error | jax, ufunc, reduce, multi-axis |
| axis_index_groups only supported for sums over just named… | exception | error | jax, psum, axis, validation |
| Convolutions with non-static strides, dilation… | exception | error | jax, jit, static-shapes, convolution, lowering |
| GMEM refs with peer ids are not supported in warpgroup… | validation | error | mosaic-gpu, pallas, prefetch, multi-gpu, peer-id |
| Non-indexing transforms on GMEM refs are not implemented. | exception | error | jax, pallas, mosaic-gpu, gmem-ref, transforms, not-implemented |
| PRNG key seed must be a scalar; got | exception | error | jax, prng, seed, scalar, shape |
| Reshape shape is not divisible by swizzle elements | validation | error | jax, pallas, reshape, swizzle, mosaic-gpu |
| scaled_matmul requires scales to match non-contract… | validation | error | jax, nn, matmul, float8, scales, shape-validation |
| Some aliased TMEM references are collective and some are… | exception | error | jax, pallas, mosaic-gpu, tmem, aliasing |
| The second argument to householder_product must not have… | validation | error | jax, linalg, householder, shape-validation |
| Cannot commute `UntilingTransform` with `ReshapeTransform`… | exception | error | jax, pallas, mosaic-gpu, reshape, tiling, layout |
| Cannot lower a pallas_call with constants. | exception | error | jax, pallas, lowering, constants, compiler |
| Currently, string arrays can only be made from NumPy… | exception | error | jax, string-dtype, type-validation |
| dtype argument to `f` must be a float dtype, got | validation | error | jax, random, f-distribution, dtype, input-validation |
| Effects not supported in partial-eval of `checkpoint`/`remat | error_code | error | pytree, jax, registry, tuple |
| Expected the same number of in_ | validation | error | jax, mosaic, mlir, layout, invariant |
| Number of Q heads must be divisible by number of KV heads… | validation | error | jax, pallas, tpu, paged-attention, gqa, head-dim-validation |
| Only single (1) and double (2) buffering are supported. Got | validation | error | jax, pallas, tpu, pipelining, buffer-count |
| re-tracing function for `jit`, but 'no_tracing' is set | exception | error | jax, config, no-tracing, cache-miss, aot |
| static_slice with partial slices does not support… | validation | error | jax, sharding, multi-device, slicing |
| tcgen05.cp only supports fully packed TMEM references | validation | error | gpu, mosaic, tcgen05, packing, tensor-memory |
| tmem_addr_ref must be in shared memory, got | exception | error | jax, mosaic, tmem, shared-memory, address-space |
| Unsupported transform | exception | error | jax, pallas, mosaic-gpu, transforms, not-implemented |
| WriteEffect should not apply to an input buffer | exception | error | jax, pallas, effects, input-mutation |
| XLA operations do not allow negative axes | validation | error | jax, lax, cumsum, negative-axis |
| arange has step == 0 | validation | error | jax, arange, zero-step, valueerror |
| {arg_ty} | exception | error | mosaic, gpu, approx, internal |
| Cannot commute unswizzle and indexer with | validation | error | jax, pallas, dtype, swizzle, indexing |
| dtype argument to `pareto` must be a float dtype, got | validation | error | jax, random, pareto, dtype, input-validation |
| expected A to be a square matrix | validation | error | jax, linalg, internal-api, shape-mismatch |
| Multimem refs are not supported in store_tiled_async | exception | error | jax, mosaic, gpu, multimem, async-store, not-implemented |
| multivariate_normal.logpdf got incompatible shapes | exception | error | jax, scipy, multivariate-normal, shape-mismatch |
| Only TiledLayouts support swizzling | exception | error | jax, mosaic-gpu, swizzle, strided-layout, store |
| points have dimension | exception | error | jax, scipy, kde, shape-validation, dimension-mismatch |
| static_slice: unrecognized index | validation | error | jax, internal, invariant |
| The context mesh should match the mesh passed to shard_map | exception | error | jax, shard-map, mesh, context-mismatch, distributed |
| The precision ' ' is not supported by dot_general on CPU | validation | error | jax, dot-general, cpu-backend, precision, fp8 |
| The unsafe_buffer_pointer() method was called on | exception | error | jax, tracing, tracer, jit, buffer-pointer |
| threefry_random_bits got invalid prng key. | exception | error | jax, prng, threefry, key-validation |
| bitcast_convert_type with different bitwidths not supported… | exception | error | jax, pallas, bitcast, dtype, not-implemented |
| Buffer poisoning is not supported on GPU yet. | exception | error | jax, pallas, buffer-poisoning, debug-flag, gpu |
| __cuda_array_interface__() is supported only for unsharded… | exception | error | jax, cuda-array-interface, cupy, sharding, interop |
| Custom JVP rule for function must produce a pair (list or… | validation | error | jax, custom-jvp, shape, dtype, autodiff |
| Gather only supports loading from VMEM, got | exception | error | jax, pallas, tpu, sparsecore, memory-space, gather |
| k must be a scalar integer; got | validation | error | jax, eye, dtype, scalar, valueerror |
| Negative padding is larger than the size of the… | validation | error | jax, convolution, padding, shape-validation |
| No 4+ dimensional dimension_number defaults. | validation | error | jax, lax, conv-transpose, dimension-numbers, high-rank |
| pargmin only accepts a single axis, got | exception | error | jax, pargmin, axis, typeerror |
| Pointwise operations on | exception | error | jax, mosaic-gpu, pointwise, quantization, int4, int8, dtype-bitwidth |
| predicate is not supported with Warpgroup lowering in… | exception | error | mosaic-gpu, pallas, jaxlib-version, predicate, tma |
| random_unwrap takes key array operand, got | exception | error | jax, prng, unwrap, dtype |
| reshape new_sizes must all be positive, got | validation | error | jax, reshape, negative-dimension, shape-validation |
| The addressable_data() method was called on | exception | error | jax, concretization, tracer, sharding, addressable-data |
| This gmm kernel only supports either (m, k) x (g, k, n) ->… | validation | error | jax, pallas, gpu, shape, moe |
| `broadcast_to` is a Triton-specific primitive. Please… | exception | error | jax, pallas, tpu, mosaic, broadcast |
| Data type bitcast is only supported for contiguous 1D… | exception | error | jax, pallas, mosaic-gpu, bitcast, stride, contiguity |
| does not accept dtype . Accepted dtypes are subtypes of . | exception | error | jax, lax, dtype-validation, type-error |
| JAX device mesh is required by multimem_load_reduce, but… | exception | error | jax, pallas, device-mesh, multimem, missing-context |
| jax.numpy.var does not yet support real dtype parameters… | validation | error | jax, numpy, variance, complex-numbers, dtype |
| Packed layout requires a GPU with at least Hopper… | exception | error | jax, cudnn, gpu-architecture, packed-attention, hopper |
| Revisited block of output in iteration . The block was… | validation | critical | jax, pallas, tpu, interpreter, overlapping-writes, grid |
| The stored value has shape | exception | error | jax, pallas, mosaic-gpu, shape-mismatch, async-store |
| Unsupported dtype for sign | exception | error | jax, pallas, tpu, dtype, sign |
| Unsupported memory space | exception | error | pallas, mosaic-gpu, memory-space, interpret-mode, not-implemented |
| Accuracy not supported | exception | error | jax, dot-general, accuracy, experimental-api |
| `block_id` must be zero when allocating a buffer for all… | exception | error | jax, pallas, mosaic-gpu, interpret-mode, allocation, block-id |
| jnp.fromiter() is not implemented because it may be… | exception | error | jax, fromiter, notimplemented, io, jit-purity |
| Kernel output in HBM but does not have trivial BlockSpec. | validation | error | jax, pallas, tpu, hbm, blockspec, interpret-mode |
| No attribute handler defined for type | exception | error | jax, mlir, custom-primitives, attributes |
| Non-slice indices are not supported in 2 minormost dims | validation | error | jax, pallas, indexing, swizzle, not-implemented |
| Only interpret mode is supported on CPU backend. | exception | error | jax, pallas, cpu, backend, interpret-mode |
| Pallas TPU requires a recent libtpu version (at least… | exception | critical | jax, pallas, tpu, libtpu, version-mismatch |
| Sharding on bias sequence dim is not allowed. | validation | error | jax, sharding, attention-bias, sequence-parallel, cudnn |
| tile with non-int block dimensions not supported yet | exception | error | jax, pallas, tile, block-spec, not-implemented |
| unreduced_psum_scatter's input cannot be varying across the… | validation | error | jax, collectives, mesh, varying, reduce-scatter |
| Unsupported transforms for | exception | error | jax, pallas, tcgen05, warp-group, transforms, not-implemented |
| Brackets are not balanced | validation | error | jax, sharding, unbalanced-parens, rule-syntax |
| Cannot broadcast shape | exception | error | jax, mosaic-gpu, broadcasting, splat-layout, shape-mismatch |
| cmp and val must have identical dtypes and shapes | validation | error | jax, pallas, triton, dtype, atomic-cas |
| Host callback lowering created too many channels. PjRt does… | exception | error | jax, host-callback, resource-limit, debugging |
| jnp.fromfile() is not implemented because it may be… | exception | error | jax, fromfile, notimplemented, io, jit-purity |
| n must be a positive integer; got | exception | error | jax, convolution-matrix, argument-validation |
| Out-of-bounds block index | validation | error | jax, pallas, tpu, blockspec, out-of-bounds |
| Permutation does not match the rank of the type ( ) | validation | error | jax, transpose, shape-mismatch |
| Scan with extensive output(s) is not supported. | validation | error | jax, pallas, triton, scan, control-flow |
| where argument of | validation | error | jax, ufunc, where-mask, numpy-compat |
| 2D+ vectors are not supported in debug_print | exception | error | gpu, mosaic, debug-print, vector, rank |
| Arrays with the splat layout can only be stored when they… | exception | error | jax, mosaic-gpu, splat-layout, warpgroup, alignment, store |
| Cannot transform type | exception | error | jax, pallas, abstract-values, transforms, type-dispatch |
| COO must have ndim=2; got | exception | error | jax, sparse, coo, shape-validation |
| could not understand directive | validation | error | jax, r-c-index-tricks, directive, numpy-compat |
| cov: dtype must be a subclass of float or complex; got | exception | error | jax, covariance, dtype-validation |
| CSR must have ndim=2; got | exception | error | jax, sparse, csr, shape-validation |
| dtype argument to `exponential` must be a float dtype, got | validation | error | jax, random, exponential, dtype, input-validation |