jax-ml/jax
Documented errors, page 2 of 23. Back to jax-ml/jax
| Code / Message | Type | Severity | Tags |
|---|---|---|---|
| The kernel function in the pallas_call | exception | error | jax, pallas, closure, constants, tracing |
| Config option already defined | exception | error | jax, config, duplicate-registration, startup |
| Effects not supported in `custom_jvp | exception | error | jax, custom-jvp, autodiff, effects |
| Number of samples without replacement | validation | error | jax, random, categorical, sampling, top-k, shape |
| Transposed memrefs are not supported in ExpandShapeOp. | validation | error | jax, mosaic-gpu, memref, expand-shape, transpose, layout-inference |
| BCSR sparse.empty: must have 2 sparse dimensions. | exception | error | jax, sparse, bcsr, shape-validation |
| shape polymorphism for Pallas does not support… | validation | error | jax, pallas, shape-polymorphism, export, experimental |
| and in op must be disjoint; got: and . | exception | error | jax, lax, gather, scatter, dimension-numbers |
| The base ref for aliases must come from a slice_smem op. | exception | error | jax, pallas, mosaic-gpu, warpgroup, smem, aliasing, not-implemented |
| The error occurred in the __reduce__ method, which may… | exception | error | jax, tracer, pickle, serialization, multiprocessing, caching |
| Unexpected dim_block_size | exception | error | jax, triton, pallas, internal, block-mapping, version-mismatch |
| use_abstract_mesh cannot change the size of the mesh. Got… | validation | error | jax, mesh, abstract-mesh, size-mismatch |
| Cannot assign layout to async load with gather indices… | validation | error | jax, mosaic, gpu, layout, alignment, async-load |
| dtypes of sorted_arr and query must match; got | exception | error | jax, dtype, searchsorted, validation |
| For flags with a corresponding contextmanager, read their… | exception | error | jax, config, contextmanager, api-migration |
| Left eigenvectors are not supported by cusolver | validation | error | jax, eig, cusolver, gpu, not-implemented |
| Can only print scalars in warp-level code. | exception | error | jax, pallas, debug-print, warp-level |
| `compute_on`'s compute_type argument must be a string. | exception | error | jax, type-error, compute-on, argument-validation |
| Value of type is not convertible to integer index. | exception | error | jax, tracer, indexing, operator-index, control-flow |
| jnp.partition for complex dtype is not implemented. | validation | error | jax, partition, complex-dtype, not-implemented, sorting |
| SMEM has shape , but expected for TMEM shape with swizzle= | validation | error | gpu, mosaic, tcgen05, shape-mismatch, swizzle, shared-memory |
| is a Unreduced -> Invariant collective. This means that the… | exception | error | jax, mesh, axis-name, collectives, sharding |
| Reverse-mode differentiation does not work for… | validation | error | jax, autodiff, grad, while-loop, fori-loop, dynamic-bounds |
| 0th dimension of leaf passed to `jax.lax.map` should be… | validation | error | jax, lax-map, sharding, multi-device, gspmd |
| dtype argument to `double_sided_maxwell` must be a float… | validation | error | jax, random, dtype-validation |
| fill_value shape cannot be broadcast to batch shape without… | exception | error | jax, nonzero, fill-value, broadcasting |
| Kernel input in HBM but does not have trivial BlockSpec. | validation | error | jax, pallas, tpu, hbm, blockspec, interpret-mode |
| Unsupported transforms | validation | error | jax, mosaic-gpu, smem-transforms, swizzle, tiling, layout-inference |
| Vector clock size ( ) must be greater than the total number… | validation | error | jax, pallas, mosaic, tpu, interpret-mode, vector-clock, config |
| COO.tree_unflatten: invalid | exception | error | jax, sparse, coo, pytree, serialization |
| custom_vjp inputs marked with nondiff_argnums must be… | exception | error | jax, custom-vjp, tracer-error, nondiff-argnums |
| Derivatives of non-symmetric eigenvectors are only valid… | validation | error | jax, autodiff, eig, gradients, numerical |
| Dimension size after padding is not at least 0, got result… | validation | error | jax, pad, negative-dimension, shape-validation |
| Nesting `compute_on` with different compute types is not… | validation | error | jax, compute-on, nesting, lowering |
| Only positional arguments are supported by debug_print on… | validation | error | jax, pallas, triton, debug-print, kwargs, gpu |
| Query and bias should have same sharding on batch and… | validation | error | jax, sharding, attention-bias, spmd, cudnn |
| async_copy requires all GMEM strides except the last one to… | validation | error | jax, mosaic-gpu, tma, memory-alignment, strides |
| dynamic_update_slice update shape must be smaller than… | exception | error | jax, lax, dynamic-update-slice, shape-mismatch |
| Integer constant out of range for i64 | exception | error | pallas, mosaic-gpu, integer-overflow, jax |
| JAX encountered invalid PRNG key data: expected… | exception | error | jax, prng, shape-mismatch, typeerror |
| When passing host local inputs to pjit, devices connected… | validation | critical | jax, mesh, multi-host, tpu, device-topology |
| all_gather_reduced is a Varying -> Reduced collective. This… | validation | error | jax, collectives, mesh, axis-name, all-gather |
| Context manager for config option requires an argument… | exception | error | jax, config, contextmanager, missing-argument |
| Inconsistent size for | validation | error | jax, einshape, einsum-notation, dimension-mismatch, shape-validation |
| matrix_transpose requires at least 2 dimensions; got | exception | error | jax, numpy, shape-validation, linalg |
| Transpose cannot be moved before a tiling transform when it… | exception | error | jax, pallas, mosaic-gpu, transpose, tiling, transforms |
| All aliased Refs must have the same memory space (SMEM or… | exception | error | jax, pallas, mosaic-gpu, smem, tmem, aliasing |
| all_gather_reduced's input cannot be reduced across the… | validation | error | jax, collectives, mesh, reduced, all-gather |
| Cannot update the mesh of the current resource environment… | validation | error | jax, mesh, axis-names, context-manager |
| Compiler params for platform | validation | error | jax, pallas, backend-mismatch, gpu, tpu |
| numpy masked arrays are not supported as direct inputs to… | validation | error | pytree, custom-node, jax, tree-flatten, api-contract |
| Value of type is not indexable. | exception | error | jax, tracer, immutable-arrays, setitem, functional-updates |
| When `check_vma=True` on `jax.shard_map`… | validation | error | jax, shard-map, vma, shapedtypestruct, mesh, tpu, config-flag |
| cannot cast from | validation | error | jax, triton, pallas, dtype, cast, gpu |
| cannot operate on split stateful generator | exception | error | jax, prng, stateful-rng, split, spawn |
| Concatenating arrays with splat layout is not supported. | exception | error | mosaic-gpu, not-implemented, splat-layout, concat, fragmented-array |
| layout attribute is only defined for TMEM refs | exception | error | jax, pallas, mosaic-gpu, tmem, gpu-kernels |
| lower and upper arguments to fori_loop must have equal… | validation | error | jax, fori-loop, dtype-mismatch, typeerror |
| MMA rhs swizzle must match lhs swizzle. | validation | error | jax, pallas, mosaic-gpu, tcgen05, mma, swizzle |
| BCSR from_scipy_sparse with nonzero n_dense/n_batch. | exception | error | jax, sparse, bcsr, scipy, not-implemented |
| Unsupported method | exception | error | jax, searchsorted, internal-api, argument-validation |
| Only single indexer is supported. | exception | error | jax, pallas, triton, atomics, indexing |
| Default value must be of type bool, got | exception | error | jax, config, typeerror, bool-default |
| Unsupported sequence length Q | exception | error | jax, cudnn, flash-attention, sequence-length, training |
| all_gather_reduced only accepts inputs that are varying. Got | validation | error | jax, collectives, mesh, spmd, all-gather |
| Collective allocations along cluster axes are not supported. | exception | error | jax, mosaic-gpu, interpret-mode, collective-allocations |
| invalid truth value for environment | validation | error | jax, environment-variable, config, bool-parsing |
| implies that array axis is partitioned times, but the… | validation | critical | jax, sharding, indivisible-shape, pjit, distributed |
| Value returned by a remat policy should be a bool or… | validation | error | jax, remat, checkpoint-policy, type-error, ad-checkpoint |
| jnp.asarray: cannot convert object of type | exception | error | jax, copy-semantics, device-placement, array-api |
| Pure callbacks do not support JVP. Please use… | exception | error | jax, autodiff, jvp, pure-callback, custom-jvp |
| Value of type is not convertible to complex. | exception | error | jax, tracer, complex-numbers, jit |
| Attempting to set log level | validation | error | jax, logging, config, valueerror |
| In MHA, expected number of 'key' heads | validation | error | jax, pallas, tpu, splash-attention, gqa, shape |
| input pinned buffers without input_output_aliases | exception | error | jax, pallas, pallas-call, input-output-alias, pinned-buffer |
| JAX encountered invalid PRNG key data: expected… | exception | error | jax, prng, dtype, uint32 |
| Non-unit strides not implemented. | exception | error | pallas, mosaic-gpu, slicing, not-implemented |
| unreduced_psum_scatter only accepts inputs that are… | validation | error | jax, collectives, mesh, spmd, reduce-scatter |
| with_memory_space_constraint only supports arrays. | validation | error | jax, pallas, abstract-value, shapedarray, type-mismatch |
| at , got fwd output type which doesn't match primal output… | exception | error | jax, custom-vjp, autodiff, type-mismatch |
| Cannot eagerly run with_memory_space_constraint. | validation | error | jax, pallas, eager-execution, primitive, memory-space |
| condition contains entries that are out of bounds | exception | error | jax, out-of-bounds, compress, length-mismatch |
| Expected to be ( ,) where `max_num_seqs` is… | validation | error | jax, pallas, tpu, paged-attention, shape-validation |
| tmem_addr_ref must be a memref or a pointer, got | exception | error | jax, mosaic, tmem, memref, argument-validation |
| Unrecognized config option | exception | error | jax, config, unknown-option, typo |
| AbstractMesh should be of the same size across all… | validation | error | jax, sharding, abstract-mesh, pjit, mesh-size-mismatch |
| Computations for m!=n are not yet supported. | exception | error | jax, scipy, legendre, not-implemented, argument-mismatch |
| end must be greater than start, but got | validation | error | jax, pallas, triton, arange, iota, empty-range, validation |
| unreduced_psum_scatter is a Unreduced -> Varying… | validation | error | jax, collectives, axis-name, mesh, reduce-scatter |
| Unsupported transforms for LHS | exception | error | jax, pallas, tcgen05, transform, not-implemented |
| accessing .backend in multi-lowering setting. This can… | exception | error | jax, lowering, multi-platform, primitive |
| Attempting to operate on barrier without indexing, but… | exception | error | jax, mosaic-gpu, barrier, indexing |
| Cannot do int indexing on TPU | validation | error | jax, pallas, tpu, indexing |
| Only scalar memrefs are supported, got | validation | error | jax, mosaic-gpu, memref, load-store, layout-inference, scalar |
| TRACER_LEAK_DEBUGGER_WARNING | console | warning | jax, tracer-leak, debugging, memory |
| Cannot pass the same ref into a mpmd map multiple times | exception | error | jax, pallas, mpmd, ref-aliasing, notimplementederror |
| Duplicate serialization registration for type | validation | error | jax, export, serialization, registry, duplicate-registration |
| Higher-order AD not supported. | exception | error | jax, pallas, tpu, splash-attention, autodiff, higher-order-ad |
| Invalid out_shape type | validation | error | jax, pallas, out-shape, shapedtypestruct, input-validation |
| PRNG key seed must be an integer; got | exception | error | jax, prng, seed, dtype, integer |