ErrLookup › jax-ml/jax

jax-ml/jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more · Python · 867 source files

Analyzed at 1e1c6a8fc0 on 2026-08-27. 2300 documented errors.

Code / MessageTypeSeverityTags
numpy masked arrays are not supported as direct inputs to JA
validation error pytree, custom-node, jax, tree-flatten, api-contract
Python int {value} too large to convert to int64
validation error pytree, custom-node, jax, tree-flatten-with-path
Python int {value} too large to convert to int32
validation error pytree, namedtuple, jax, tree-flatten
The names should be exclusive and should not intersect in `n
validation error pytree, namedtuple, jax, pickle, tree-flatten-with-path
The return value of the policies should be a boolean. Got: {
validation warning pytree, deprecation, sets, generators, jax
prevent_cse must be a bool or tuple of bools, got {type(prev
validation error pytree, unflatten, jax, leaf-count
the `static_argnums` argument to `jax.checkpoint` / `jax.rem
validation error pytree, unflatten, jax, leaf-count
the `static_argnums` argument to `jax.checkpoint` / `jax.rem
validation error pytree, tree-map, structure-mismatch, jax
Effects not supported in AD of `checkpoint`/`remat`: {disall
error_code error pytree, none-handling, jax, tree-map, breaking-change
non-bool prevent_cse together with static_argnums/static_arg
error_code error pytree, jax, structure-mismatch, from-iterable
the rematted computation's closure contains a mutable array
validation error pytree, jax, registry, compose
Effects not supported in partial-eval of `checkpoint`/`remat
error_code error pytree, jax, registry, tuple
stop_gradient only works on valid JAX arrays, but input argu
validation error pytree, jax, pickle, version-mismatch
{self.__class__.__name__} has no attribute {name}
error_code error pytree, jax, pickle, malformed-data
Mapped away dimension of inputs passed to vmap should be sha
validation error jax, vmap, sharding, mesh, distributed
Unmapped values passed to vmap cannot be sharded along the m
validation error jax, vmap, sharding, mesh
{name} wrapped function must be passed at least one argument
validation error jax, vmap, axis-size, api-misuse
{name} was requested to map a value of non-array type {core.
validation error jax, vmap, in-axes, type-error
{name} was requested to map its argument along axis {axis},
validation error jax, vmap, shape, rank, in-axes
{name} must have at least one non-None value in in_axes or a
validation error jax, vmap, in-axes, axis-size
{''.join(msg)[:-2]}
validation error jax, vmap, shape, batch-size
primal and tangent arguments to jax.jvp must be tuples or li
validation error jax, jvp, autodiff, api-misuse, type-error
primal and tangent arguments to jax.jvp must have the same t
validation error jax, jvp, pytree, autodiff
primal and tangent arguments to jax.jvp do not match; dtypes
validation error jax, jvp, dtype, autodiff, float0
jvp called with different primal and tangent shapes;Got prim
validation error jax, jvp, shape, autodiff
expected {in_avals.tree}, got {tangents_ft.tree}
validation error jax, linearize, jvp, pytree
linearized function called on tangent values inconsistent wi
validation error jax, linearize, tangent, aval, mixed-precision
reduce_axes argument to vjp is deprecated
validation error jax, vjp, deprecated, migration
unexpected JAX type (e.g. shape/dtype) for gradient ref pass
validation error jax, vjp, refs, gradient-accumulation, dtype
the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
validation error jax, vjp, refs, autodiff
{_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
validation error jax, vjp, refs, autodiff
But the tree structures differ:
validation error jax, vjp, pytree, refs
the VJP function was applied before restoring its not-saveab
validation error jax, vjp, memory, checkpointing, saveable-args
Accumulator aval mismatch: expected {aval}, got {acc.aval}
validation error jax, vjp, accumulator, internal-api, dtype
{full_name} must be a pytree prefix with bool leaves or a tu
validation error jax, pytree, validation, bool-prefix
{full_name} must form a tree prefix of the corresponding val
validation error jax, pytree, tree-prefix
{full_name} must form a tree prefix of the corresponding val
validation error jax, pytree, arity
structure of the differentiated function {jaxpr.debug_info.f
validation error jax, vjp, autodiff, pytree
unexpected JAX type (e.g. shape/dtype) for argument to VJP f
validation error jax, vjp, dtype, shape-mismatch
{_vjp_too_many_args(name, len(extra_args) + 1)}
validation error jax, vjp, arity, argument-count
reduce_axes argument to transpose is deprecated
validation error jax, linear-transpose, deprecated-api
linear_transpose only supports [float or complex] -> [float
validation error jax, linear-transpose, dtype
cotangent tree does not match function output, expected {out
validation error jax, linear-transpose, pytree
cotangent type does not match function output, expected {out
validation error jax, linear-transpose, dtype, shape-mismatch
iteration over a 0-d array
exception error jax, numpy-interop, iteration, scalar, typeerror
to_dlpack can only pack a dlpack tensor from an array on a s
exception error jax, dlpack, interop, sharding, multi-device
__dlpack__ only supported for unsharded arrays.
exception error jax, dlpack, sharding, buffer-protocol, interop
Unknown GPU platform for __dlpack__: {platform_version}
exception error jax, dlpack, gpu, platform-detection, version-mismatch
Couldn't get local_hardware_id for __dlpack__
exception error jax, dlpack, gpu, device-init, environment
__dlpack__ device only supported for TPU pinned host memory
exception error jax, dlpack, tpu, pinned-memory, interop
__dlpack__ device only supported for CPU, GPU and TPU pinned
exception error jax, dlpack, backend, unsupported-platform, interop
unsafe_buffer_pointer() is supported only for unsharded arra
exception error jax, cuda, raw-pointer, sharding, unsupported-operation
__cuda_array_interface__() is supported only for unsharded a
exception error jax, cuda-array-interface, cupy, sharding, interop
Buffer callbacks do not support transpose. Please use `jax.c
exception error jax, autodiff, gradient, callback, buffer-callback
multi-platform lowering for buffer_callback
exception error jax, xla, multi-platform, lowering, export
`buffer_callback` not supported on {platform} backend.
exception error jax, backend, tpu, unsupported-platform, buffer-callback
buffer_callback callback must not return any values.
exception error jax, callback, host-callback, return-value, buffer-callback
jax.pure_callback failed to find a local CPU device to place
exception error jax, environment, pure-callback, cpu-device, jax-platforms
Pure callbacks do not support JVP. Please use `jax.custom_jv
exception error jax, autodiff, jvp, pure-callback, custom-jvp
Pure callbacks do not support transpose. Please use `jax.cus
exception error jax, autodiff, vjp, transpose, pure-callback, custom-vjp
callbacks are only supported in spmd computations when all m
exception error jax, spmd, sharding, mesh, pure-callback, manual-axes
callbacks do not support specifying sharding inside spmd com
exception error jax, spmd, sharding, pure-callback, manual-axes
pure_callback only supports SingleDeviceSharding, but got {t
exception error jax, sharding, pure-callback, single-device-sharding
Partitioned callback not implemented on {platform} backend.
exception error jax, callback, tpu, sharding, backend-support
Partitioned callback not supported with return values.
exception error jax, callback, sharding, api-contract
Mismatched number of outputs from callback. Expected: {}, Ac
exception error jax, callback, runtime-validation, shape-mismatch
Incorrect output shape for return value #{i}: Expected: {out
exception error jax, callback, shape-mismatch, runtime-validation
Incorrect output dtype for return value #{i}: Expected: {out
exception error jax, callback, dtype-mismatch, runtime-validation
{str(exc)}
exception error jax, checkify, assertion, debugging
Checkify does not support batched while-loops (checkify-of-v
exception error jax, checkify, vmap, while-loop, transformation-order
Mesh must be provided for shard_map with checkify.
exception error jax, checkify, shard-map, mesh, sharding
Unsupported aval type: {type(v)}
exception error jax, checkify, shard-map, custom-aval, extensibility
{prim_name} takes a scalar pred as argument, got {pred}
exception error jax, checkify, type-error, scalar-predicate
Formatting arguments to checkify.check need to be PyTrees of
exception error jax, checkify, typeerror, assertions, jit
check_error takes an Error as argument, got type {type(error
exception error jax, checkify, typeerror, api-misuse
`compute_on`'s compute_type argument must be a string.
exception error jax, type-error, compute-on, argument-validation
Nesting `compute_on` with different compute types is not all
validation error jax, compute-on, nesting, lowering
invalid truth value {val!r} for environment {varname!r}
validation error jax, environment-variable, config, bool-parsing
Unrecognized config option: {name}
exception error jax, config, unknown-option, typo
For flags with a corresponding contextmanager, read their va
exception error jax, config, contextmanager, api-migration
Config option {name} already defined
exception error jax, config, duplicate-registration, startup
bool() not supported for instances of type '{0}' (did you me
exception error jax, config, typeerror, truthiness
Context manager for {state.__name__} config option requires
exception error jax, config, contextmanager, missing-argument
Default value must be of type bool, got {default} of type {g
exception error jax, config, typeerror, bool-default
Default value must be of type str, got {default} of type {ge
exception error jax, config, typeerror, enum
Invalid value "{default}" for JAX flag {name}
validation error jax, config, valueerror, environment-variable, enum
new enum value must be in {enum_values}, got {new_val} of ty
validation error jax, config, valueerror, enum
Default value must be of type str or None, got {default} of
exception error jax, config, typeerror, enum
new enum value must be None or in {enum_values}, got {new_va
validation error jax, config, valueerror, enum
Default value must be of type {enum_class}, got {default} of
exception error jax, config, typeerror, enum
Invalid value "{default_str}" for JAX flag {name}
validation error jax, config, valueerror, environment-variable, enum
new enum value must be an instance of {enum_class}, got {new
validation error jax, config, typeerror, enum
Default value must be of type int, got {default} of type {ge
exception error jax, config, typeerror, int
Invalid value "{default_env}" for JAX flag {name}
validation error jax, config, valueerror, environment-variable, int
new int config value must be None or of type int, got {new_v
validation error jax, config, valueerror, int
Default value must be of type float, got {default} of type {
exception error jax, config, typeerror, float
new float config value must be None or of type float, got {n
validation error jax, config, valueerror, float
new string config value must be of type str, got {new_val} o
validation error jax, config, typeerror, string
new string config value must be None or of type str, got {ne
validation error jax, config, valueerror, string
Value of type {type(self)} is not convertible to float.
exception error jax, tracer, type-conversion, jit, float
Value of type {type(self)} is not convertible to complex.
exception error jax, tracer, complex-numbers, jit
Value of type {type(self)} is not convertible to hex.
exception error jax, tracer, hex, formatting, debugging
Value of type {type(self)} is not convertible to oct.
exception error jax, tracer, oct, formatting
Value of type {type(self)} is not convertible to integer ind
exception error jax, tracer, indexing, operator-index, control-flow
The error occurred in the __reduce__ method, which may indic
exception error jax, tracer, pickle, serialization, multiprocessing, caching
Value of type {type(self)} is not indexable.
exception error jax, tracer, immutable-arrays, setitem, functional-updates
Value of type {type(self)} is not compatible with the Array
exception error jax, tracer, numpy-interop, array-api, nep-35
The '{name}' method is not available on {self._error_repr()}
exception error jax, tracer, block-until-ready, async, attribute-access
The 'sharding' attribute is not available on {self._error_re
exception error jax, tracer, sharding, attribute-access, distributed
{self.__class__.__name__} has no attribute {name}
exception error jax, tracer, attributeerror, jit, duck-typing
The addressable_data() method was called on {self._error_rep
exception error jax, concretization, tracer, sharding, addressable-data
The delete() method was called on {self._error_repr()}.{self
exception error jax, concretization, tracer, memory-management, delete
The devices() method was called on {self._error_repr()}.{sel
exception error jax, concretization, tracer, devices, placement
The global_shards property was called on {self._error_repr()
exception error jax, concretization, tracer, sharding, global-shards
The is_deleted() method was called on {self._error_repr()}.{
exception error jax, concretization, tracer, memory-management, is-deleted
The is_fully_addressable property was called on {self._error
exception error jax, concretization, tracer, sharding, distributed
The is_fully_replicated property was called on {self._error_
exception error jax, concretization, tracer, sharding, replication
The on_device_size_in_bytes() method was called on {self._er
exception error jax, concretization, tracer, memory, profiling
The traceback property was called on {self._error_repr()}.{s
exception error jax, concretization, tracer, traceback, debugging
The unsafe_buffer_pointer() method was called on {self._erro
exception error jax, tracing, tracer, jit, buffer-pointer
unbound axis name: {axis_name}
exception error jax, axis-name, vmap, pmap, collectives
Triggering __jax_array__() during abstractification is no lo
exception error jax, pytree, jax-array, custom-class, tracing
Cannot interpret value of type {typ} as an abstract array; i
exception error jax, type-error, abstract-value, tracing, dtype
Argument '{x}' of type '{typ}' is not a valid JAX type
exception error jax, type-error, jit, argument-validation
TracerBoolConversionError
exception error jax, tracer, bool-conversion, control-flow, jit
TracerIntegerConversionError
exception error jax, tracer, integer-conversion, indexing, jit
The problem arose with the `{fname}` function.
exception error jax, tracer, concretization, jit, float-conversion
Sharding spec {spec} implies that array axis {dim} is partit
exception error jax, sharding, named-sharding, mesh, partitionspec
Length of sharding.spec ({len(out_s.spec)}) must be equal to
exception error jax, sharding, partition-spec, shape-mismatch
Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of
exception error jax, sharding, mesh, abstract-mesh, type-mismatch
sharding with memory_kind is not allowed. Please use `jax.de
exception error jax, sharding, memory-kind, device-put
Axes mentioned in `manual_axis_type` field of ShapedArray sh
exception error jax, sharding, manual-axis-type, mesh, validation
varying and unreduced cannot have common mesh axes. Got vary
validation error jax, sharding, manual-axis-type, validation
varying and reduced cannot have common mesh axes. Got varyin
validation error jax, sharding, manual-axis-type, validation
Expected unreduced_kind to be of type `jax.sharding.Unreduce
validation error jax, sharding, unreduced-kind, type-error
`unreduced_kind` should be `None` when `unreduced` is an emp
validation error jax, sharding, unreduced-kind, validation
{name} argument of ManualAxisType should of type `frozenset`
validation error jax, sharding, type-error, constructor-validation
The layout of ShapedArray should not be `AutoLayout` when la
validation error jax, layout, auto-layout, configuration
len() of unsized object
exception error jax, numpy, len, zero-d-array
{name} cannot accept args which are unreduced. Got {a.str_sh
validation error jax, sharding, collectives, unreduced, spmd
{name} cannot accept args with unreduced_kind={a.mat.unreduc
validation error jax, sharding, collectives, unreduced-kind, spmd
{name} cannot accept args which are reduced. Got {a.str_shor
validation error jax, sharding, collectives, reduced, spmd
Primitive {prim_name} requires varying manual axes to match,
validation error jax, shard-map, spmd, internal-error, workaround
array ref with memory space only works inside of a `jit`.
exception error jax, refs, memory-space, jit, not-implemented
pinned array ref only works inside of a `jit`.
exception error jax, refs, pinned-memory, jit, not-implemented
{self!r} has no `sharding`.
exception error jax, aval, attribute-error, tokens
Unsupported sequence length Q {T}, KV {S}.
exception error jax, cudnn, flash-attention, sequence-length, training
Packed layout requires a GPU with at least Hopper architectu
exception error jax, cudnn, gpu-architecture, packed-attention, hopper
mla requires cudnn version >= 9.10 and at least hopper arch.
exception error jax, cudnn, mla, version-mismatch, gpu-architecture
cuDNN is not detected.
exception critical jax, cudnn, installation, environment, missing-dependency
Currently only support batch_dim in [0, None], but got {dim=
exception error jax, vmap, cudnn, batching, in-axes
Query, key and value should have same sharding.
validation error jax, sharding, gspmd, multi-device, cudnn
Sharding on sequence dim is not allowed.
validation error jax, sharding, sequence-parallel, spmd, cudnn
Sharding on head dim is not allowed.
validation error jax, sharding, tensor-parallel, spmd, cudnn
Query and bias should have same sharding on batch and num_he
validation error jax, sharding, attention-bias, spmd, cudnn
Sharding on bias sequence dim is not allowed.
validation error jax, sharding, attention-bias, sequence-parallel, cudnn
Require q_seqlen and kv_seqlen to use packed layout
validation error jax, attention, varlen, packed-layout, cudnn
shardings should container 4 inputs, but got {len(shardings)
validation error jax, fp8, sharding, spmd, matmul
shardings specs rank should be 3, but got lhs: {len(lhs.spec
validation error jax, fp8, sharding, rank-mismatch, matmul
shardings spec for batch dim should be same, but got lhs: {l
validation error jax, fp8, sharding, batch-dim, matmul
Unrecognized mode: {config.mode}.
validation error jax, fp8, quantization, config, cudnn
Only support preferred_element_type in (f32, bf16, f16), but
validation error jax, dtype, fp8, matmul, unsupported-type
The input arguments to the custom_vmap-decorated function {d
validation error jax, custom-vmap, kwargs, argument-binding
No batching rule defined for custom_vmap function {debug_fun
exception error jax, custom-vmap, vmap, missing-rule
structure of output value and output batching specification
validation error jax, custom-vmap, pytree, vmap
structure of output returned by custom vmap rule ({rule_name
validation error jax, custom-vmap, pytree, vmap
Can't use ``defjvps`` with ``nondiff_argnums``.
validation error jax, custom-jvp, autodiff, api-misuse
No JVP defined for custom_jvp function {primal_name} using d
exception error jax, custom-jvp, autodiff, missing-rule
The input arguments to the custom_jvp-decorated function {pr
validation error jax, custom-jvp, kwargs, argument-binding
Custom JVP rule {jvp_name} for function {primal_name} must p
validation error jax, custom-jvp, autodiff, return-shape
Custom JVP rule {jvp_name} for function {primal_name} must p
validation error jax, custom-jvp, pytree, autodiff
Custom JVP rule {jvp_name} for function {primal_name} must p
validation error jax, custom-jvp, pytree, autodiff
Custom JVP rule {jvp_name} for function {primal_name} must p
validation error jax, custom-jvp, shape, dtype, autodiff
Custom JVP rule must produce primal and tangent outputs with
validation error jax, custom-jvp, tangent, shape, dtype
Custom JVP rule must produce primal and tangent outputs with
validation error jax, custom-jvp, tangent, pytree, autodiff
Effects not supported in `custom_jvp`: {disallowed_effects}
exception error jax, custom-jvp, effects, purity, jit
custom_gradient function used with with_logs=True must retur
exception error jax, custom-gradient, typeerror, vjp
The inputs to the closure produced by closure_convert must h
exception error jax, closure-convert, pytree, typeerror
transpose output pytree structure must match that of linear
exception error jax, linear-call, transpose, pytree
remat optimization for custom_vjp does not support forward f
exception error jax, custom-vjp, remat, side-effects
remat optimization for custom_vjp does not support higher-or
exception error jax, custom-vjp, remat, higher-order-ad
keyword arguments could not be resolved to positions
exception error jax, custom-partitioning, kwargs, typeerror
Custom Partitioning rules must return Sharding.
exception error jax, custom-partitioning, sharding, tpu
Mismatch in result shapes. %s vs %s
exception error jax, custom-partitioning, sharding, shape-mismatch
Unknown keyword arguments: {sharding_rule_dict}
exception error jax, custom-partitioning, shardy, sharding-rule
Shardy is used, but sharding propagation callbacks instead o
exception error jax, shardy, custom-partitioning, migration, tpu
Custom-partitioned function {function!r} does not support GS
exception error jax, gspmd, shardy, custom-partitioning, config
sharding_rule callable must produce either an SdyShardingRul
exception error jax, shardy, sharding-rule, custom-partitioning
Factor names have to start with a letter, but got '{factor[0
exception error jax, shardy, sharding-rule, validation
Unknown character '{char}'
exception error jax, shardy, sharding-rule, validation
A compound factor should contain at least two factors
exception error jax, shardy, sharding-rule, validation
Each element of CompoundFactor must be a str, but got {type(
exception error jax, shardy, sharding-rule, typeerror
Ellipsis can't be used in a compound factor
exception error jax, shardy, sharding-rule, ellipsis
Each element of ArrayMapping must be a str or CompoundFactor
exception error jax, shardy, sharding-rule, typeerror
Ellipsis can only be used at the beginning of a dimension
exception error jax, shardy, sharding-rule, ellipsis
Factor {factor} is not used in the rule, but size is provide
exception error jax, sharding, factor-sizes, validation
Factor {factor} is only used in compound factors; must speci
exception error jax, sharding, compound-factor, factor-sizes
Factor {factor} represents a whole dimension; do not specify
exception error jax, sharding, factor-sizes, redundant-config
{kind} must be a tuple of factors
exception error jax, sharding, type-error, tuple
{kind} contains duplicated factors
validation error jax, sharding, duplicate-values, validation
Factor {factor} in {kind} is not used in the rule
validation error jax, sharding, unknown-factor, validation
Factor {factor} can only be in one of the reduction, need re
validation error jax, sharding, conflicting-config, validation
Compound factors should be one level, nested brackets are no
validation error jax, sharding, compound-factor, rule-syntax
Brackets are not balanced
validation error jax, sharding, unbalanced-parens, rule-syntax
Brackets should contain at least two factors
validation error jax, sharding, compound-factor, rule-syntax
Factor names have to start with a letter, but got '{char}'
validation error jax, sharding, naming, rule-syntax
Brackets are not balanced in rule: '{rule}'
validation error jax, sharding, unbalanced-parens, rule-syntax
rule must be a str, but got {type(rule)}
validation error jax, sharding, type-error, rule-syntax
factor_sizes must be a dict of str to int, but got {factor_s
validation error jax, sharding, type-error, factor-sizes
Unknown character '{BATCHING}'
validation error jax, sharding, reserved-character, rule-syntax
Character '.' must be used inside ellipsis '...'
validation error jax, sharding, ellipsis, rule-syntax
There is no -> in rule: '{rule}'
validation error jax, sharding, missing-separator, rule-syntax
Sharding rule has {len(rule.operand_mappings)} operands, but
validation error jax, sharding, custom-partitioning, mlir, sdy
Argument '{arg}' of type {type(arg)} is not a valid JAX type
validation error jax, type-validation, argument-validation
nan
exception error jax, nan, numerical-debugging
inf
exception error jax, inf, numerical-overflow
For a cross-host reshard in multi-controller JAX, input and
validation error jax, multi-controller, resharding, device-mismatch
mTLS for the JAX distributed service requires jaxlib 0.11.2
exception error jax, distributed, mtls, version-mismatch
coordinator_address should be defined.
validation error jax, distributed, coordinator, config
Number of processes must be defined.
validation error jax, distributed, world-size, config
The process id of the current process must be defined.
validation error jax, distributed, rank, config
process_id must be a nonnegative int. Got process_id={proces
validation error jax, distributed, type-error, rank
num_processes must be a positive int. Got num_processes={num
validation error jax, distributed, type-error, world-size
process_id and num_processes must be nonnegative, with proce
validation error jax, distributed, rank-validation
coordinator_bind_address should be defined.
validation error jax, distributed, network, bind-address
distributed.initialize should only be called once.
exception error jax, distributed, double-init, lifecycle
Preemption sync manager should only be initialized once.
exception error jax, distributed, preemption, singleton
jax.distributed.initialize() must be called before any JAX c
exception critical jax, distributed, xla-backend, ordering
Expected a string or dtype-like object; got {dtype=}
validation error jax, dlpack, dtype, none-guard
Specified {device=} which requires a copy since the source d
validation error jax, dlpack, device-mismatch, copy
Argument to to_dlpack must be a jax.Array, got {type(x)}
validation error jax, dlpack, type-check, typeguard
The device specification passed to to_dlpack contains an uns
exception error jax, dlpack, device-type, buffer-protocol
JAX does not support any version below {MIN_DLPACK_VERSION}
exception error jax, dlpack, version-negotiation, protocol
Specified {device=} which requires a copy since the source d
validation error jax, dlpack, device-mismatch, copy
from_dlpack can only unpack a dlpack tensor onto a singular
validation error jax, dlpack, sharding, multi-device
The array passed to from_dlpack must have __dlpack__ and __d
validation error jax, dlpack, type-check, buffer-protocol
Array passed to from_dlpack is on unsupported device type (D
validation error jax, dlpack, device-type, interop
Specified input which requires a copy since the source data
validation error jax, dlpack, alignment, zero-copy
JAX requires ml_dtypes version 0.5 or newer; installed versi
validation critical jax, dependency, version, ml-dtypes
dtype cannot be None.
validation error jax, dtype, none-guard, validation
unexpected input: {dtype=}
validation error jax, dtype, validation
dtype {dtype!r} not understood
validation error jax, dtype, type-check, canonicalization
Invalid scalar value {x}
validation error jax, scalar, type-check
Python int {value} too large to convert to {dtype}
validation error jax, overflow, int64, x64
Unrecognized {kind=} expected one of {list(_dtype_kinds.keys
validation warning jax, dtype, array-api, validation
Expected kind to be a dtype, string, or tuple; got {kind=}
validation error jax, dtype, isdtype, argument-validation
{dtype=} is not a valid dtype for JAX type promotion.
validation error jax, type-promotion, dtype, lattice
Input dtypes {tuple(str(n) for n in nodes)} have no availabl
validation error jax, type-promotion, casting, dtype
Dtype {dtype} is not a valid JAX array type. Only arrays of
validation error jax, dtype, validation, numeric-types
Explicitly requested dtype {}{} is not available. To enable
validation error jax, x64, dtype, configuration
Invalid argument to dtype: {x}.
validation error jax, dtype, none-argument
Value '{x}' with dtype {dt} is not a valid JAX array type. O
validation error jax, dtype, string-dtype, validation
Cannot determine dtype of {x}
validation error jax, dtype, array-like, inference-failure
at least one array or dtype is required
validation error jax, result-type, argument-validation
dtype must be specified.
validation error jax, dtype, argument-validation
Passing an array as a dtype argument is no longer supported;
validation error jax, dtype, api-removal, migration
JAX only supports number, bool, and string dtypes, got dtype
validation error jax, dtype, validation, numeric-types
iteration over a 0-d array
validation error jax, earray, iteration, zero-d-array
len() of unsized object
validation error jax, earray, len, zero-d-array
The error code state and the predicate must be on the same m
validation error jax, error-checking, sharding, mesh, distributed
raise_if_error() should not be called within a traced contex
validation error jax, error-checking, jit, tracer
No VJP is available
validation error jax, export, vjp, serialization
Duplicate serialization registration for type `{nodetype}`.
validation error jax, export, serialization, registry, duplicate-registration
Duplicate serialization registration for serialized_name `{s
validation error jax, export, serialization, registry, name-collision
serializing DisabledSafetyCheck: {check}
exception error jax, export, serialization, flatbuffers
Cannot deserialize DisabledSafetyCheck with unknown kind: {k
exception error jax, export, deserialization, version-mismatch
Encountered dimension variable '{self.var}' that is not appe
exception error jax, export, shape-polymorphism, dimension-variables
Cannot divide {self} by {divisor}.
exception error jax, shape-polymorphism, floor-division, export
__pow__ modulo not implemented
exception error jax, shape-polymorphism, pow, not-implemented
Symbolic dimension cannot be raised to non-integer powers: '
exception error jax, shape-polymorphism, pow, validation
Symbolic dimension cannot be raised to negative powers: '{se
exception error jax, shape-polymorphism, pow, negative-exponent
__rpow__ modulo not implemented
exception error jax, shape-polymorphism, rpow, not-implemented
Symbolic dimension '{self}' used in a context that requires
exception error jax, shape-polymorphism, int-conversion, tracing
The symbolic constraints should be a sequence of strings. Go
exception error jax, shape-polymorphism, constraints, api-misuse
SymbolicScope constraint must be a string: got {repr(c_str)}
exception error jax, shape-polymorphism, constraints, type-validation
Constraint parsing error: must contain one of '==' or '>=' o
exception error jax, shape-polymorphism, constraints, parsing
Unsatisfiable explicit constraint: {constr.debug_str}
exception error jax, shape-polymorphism, constraints, unsatisfiable
Invalid equality constraint: {e1} == {e2}. The left-hand-sid
exception error jax, shape-polymorphism, constraints, normalization
Found multiple equality constraints with the same left-hand-
exception error jax, shape-polymorphism, constraints, duplicate-key
Invalid mixing of symbolic scopes {when}.\nExpected {self_de
exception error jax, shape-polymorphism, scope, export
Symbolic dimension comparison {cmp_str()} is inconclusive.{d
exception error jax, shape-polymorphism, comparison, inconclusive
Symbolic dimension {operation_name} not supported for {p}.
exception error jax, shape-polymorphism, type-error, coercion
Encountered unexpected shape dimension {d}
exception error jax, export, shape-polymorphism, type-error
for grad support, subclass {type(self)} must implement `vjp_
validation error jax, autodiff, custom-primitive, not-implemented
for jvp support, subclass {type(self)} must implement `jvp`
validation error jax, autodiff, forward-mode, not-implemented
for linearize support, subclass {type(self)} must implement
validation error jax, autodiff, linearize, not-implemented
for transpose support, subclass {type(self)} must implement
validation error jax, autodiff, transpose, not-implemented
for vmap support, subclass {type(self)} must implement `batc
validation error jax, vmap, batching, custom-primitive
under vmap, the {member_name} of {name} produced an output b
validation error jax, vmap, custom-vjp, custom-jvp, batching
input type mismatch for {_prim}
validation error jax, type-mismatch, tracing, custom-primitive
{type(_prim).__name__} returned structured residuals from `v
validation error jax, autodiff, vjp, api-contract
tree mismatch during linearization of {prim=}. Expected: {pr
error_code error jax, pytree, linearize, autodiff
{type(_prim).__name__}.vjp_bwd should return None or a dict
validation error jax, autodiff, vjp, api-contract
{type(_prim).__name__}.transpose should return None or a dic
validation error jax, autodiff, transpose, api-contract
subclass {type(self)} can't set both `jvp = jvp_from_lin` an
validation error jax, autodiff, circular-definition, misconfiguration
`linearize_from_jvp` is a pair of rules, not a single rule;
validation error jax, api-misuse, unpacking
`vjp_from_jvp` is a pair of rules, not a single rule; unpack
validation error jax, api-misuse, unpacking
`vjp_from_lin` is a pair of rules, not a single rule; unpack
validation error jax, api-misuse, unpacking
Custom VJP bwd rule {self.bwd} was registered with defvjp_wi
exception error jax, custom-vjp, api-contract, logging
Custom VJP bwd rule {self.bwd} was registered with defvjp_wi
exception error jax, custom-vjp, logging, type-validation
Custom VJP bwd rule {self.bwd} must produce a tuple but got
exception error jax, custom-vjp, return-type
Custom VJP bwd rule {self.bwd} must produce a tuple of lengt
exception error jax, custom-vjp, arity-mismatch
Effects not supported in `custom_jvp`: {disallowed}
exception error jax, custom-jvp, autodiff, effects
at {keystr(path)}, got fwd output type {ty.str_short()} whic
exception error jax, custom-vjp, autodiff, type-mismatch
{result}the bwd rule attached to {primal_sourceinfo} produce
exception error jax, custom-vjp, autodiff, type-mismatch
No VJP defined for custom_vjp function {self.f.__name__} usi
exception error jax, custom-vjp, missing-rule, autodiff
custom_vjp inputs marked with nondiff_argnums must be static
exception error jax, custom-vjp, tracer-error, nondiff-argnums
Custom JVP rule {jvp_name} for function {self.traced.fun_nam
exception error jax, custom-jvp, autodiff, return-contract
Custom JVP rule {jvp_name} for function {self.traced.fun_nam
exception error jax, custom-jvp, pytree-mismatch
Custom JVP rule {jvp_name} for function {self.traced.fun_nam
exception error jax, custom-jvp, pytree-mismatch
Custom JVP rule must produce primal and tangent outputs with
exception error jax, custom-jvp, tangent-mismatch
Custom JVP rule must produce primal and tangent outputs with
exception error jax, custom-jvp, tangent-mismatch, multi-output
Can't use ``defjvps`` with ``nondiff_argnums``.
exception error jax, custom-jvp, defjvps, api-misuse
No JVP defined for custom_jvp function {self.f.__name__} usi
exception error jax, custom-jvp, missing-rule, autodiff
The input arguments to the custom_jvp-decorated function {se
exception error jax, custom-jvp, kwargs, signature-binding
custom_jvp inputs marked with nondiff_argnums must be static
exception error jax, custom-jvp, tracer-error, nondiff-argnums
custom_jvp-decorated function {self.f} closed over a {type(t
exception error jax, custom-jvp, tracer-error, closure
Sharded hypothesis runner does not support `data()` inside `
exception error jax, testing, hypothesis, sharding
Unknown resize method "{s}"
validation error jax, image-resize, invalid-argument
shape must have length equal to the number of dimensions of
validation error jax, image-resize, shape-mismatch
No constant handler for type: {type(val)}
exception error jax, constants, type-error, tracing
numpy masked arrays are not supported as direct inputs to JA
exception error numpy, jax, masked-array, input-validation
A ShapeDtypeStruct does not have a value and cannot be used
exception error jax, shape-dtype-struct, constants
Unsupported scalar attribute type: {type(val)}
exception error jax, dtype, mlir, numpy
NumPy arrays with zero strides are not supported as MLIR att
exception error numpy, strides, jax, mlir
No attribute handler defined for type: {type(val)}
exception error jax, mlir, custom-primitives, attributes
accessing .backend in multi-lowering setting. This can occur
exception error jax, lowering, multi-platform, primitive
the platform for the specified backend {xb.canonicalize_plat
exception error jax, platform-mismatch, lowering, backend
Host callback lowering created too many channels. PjRt does
exception error jax, host-callback, resource-limit, debugging
Invalid compute type {c_type}. Current supported values are
exception error jax, mlir, input-validation, enum, compute-type
Cannot lower jaxpr with effects: {closed_jaxpr.effects}
exception error jax, effects, custom-partitioning, lowering, pmap
custom_call backend_config unexpected type: {backend_config}
exception error jax, mlir, custom-call, backend-config, type-error
Error refining shapes. {dump_module_message(module, "before_
exception error jax, polymorphic-shapes, dynamic-shapes, mlir, lowering
Cannot use `partial_eval_jaxpr_custom` with stateful jaxprs.
validation error jax, remat, partial-eval, stateful-jaxpr, control-flow
Value returned by a remat policy should be a bool or `ad_che
validation error jax, remat, checkpoint-policy, type-error, ad-checkpoint
`JaxprInputEffect` {eff} does not have a corresponding jaxpr
validation error jax, jaxpr, effects, internal-invariant, mutable-arrays
{primitive}.abstract_eval() method should return a tuple or
validation error jax, custom-primitives, abstract-eval, multiple-results, api-contract
re-tracing function {fun} for `jit`, but 'no_tracing' is set
exception error jax, config, no-tracing, cache-miss, aot
function {dbg.func_src_info} traced for {dbg.traced_for} ret
validation error jax, tracing, type-error, return-value, jit
function {dbg.func_src_info} traced for {dbg.traced_for} ret
validation error jax, mutable-arrays, refs, tracing, return-value
numpy masked arrays are not supported as direct inputs to JA
validation error jax, numpy, masked-array, input-validation, sharding
No pxla_result_handler for type: {type(aval)}
error_code error jax, pxla, custom-aval, result-handler, extension-api
JAX tried to execute function {self.name}, but the no_execut
exception error jax, config, no-execution, environment, diagnostics
unbound axis name: {axis_name}
validation error jax, mesh, collectives, axis-name, sharding
AbstractMesh should be of the same size across all shardings
validation error jax, sharding, abstract-mesh, pjit, mesh-size-mismatch
AbstractMesh size: {abstract_mesh.size} does not match the d
validation error jax, sharding, abstract-mesh, device-assignment, pjit
The following ordered effects are not supported for more tha
validation error jax, effects, multi-device, ordered-effects, pjit, debug-print
Passing lowering_platforms via jax.export or jit(f).trace(*a
validation error jax, abstract-mesh, lowering-platforms, export, aot, pjit
State effect not supported in vmap-of-cond.
exception error jax, vmap, cond, state-effects, experimental
IO effect not supported in vmap-of-cond.
exception error jax, vmap, cond, io-callback, host-callback
State effect not supported in cond partial-eval.
exception error jax, cond, partial-eval, state-effects, autodiff
lax.platform_dependent: the '{pname}' branch must be a calla
validation error jax, lax, platform-dependent, typeerror, callable
Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.
validation error jax, lax, platform-dependent, gpu, invalid-argument
lax.platform_dependent: the 'default' branch must be a calla
validation error jax, lax, platform-dependent, typeerror, callable
Effects not supported in `scan`: {disallowed_effects}
exception error jax, scan, effects, scan3, experimental
`unroll` must be a `bool` or a non-negative `int`.
validation error jax, scan, unroll, valueerror, argument-validation
lax.scan: f argument should be a callable.
validation error jax, scan, typeerror, callable, api-misuse
zero-length scan is not supported in disable_jit() mode beca
validation error jax, scan, disable-jit, zero-length, edge-case
scan body output must be a pair, got {}.
validation error jax, scan, typeerror, return-shape, pytree
must provide `length` to `scan`, since the leading-axis size
validation error jax, scan, hijax, length, experimental
scan got value with no leading axis to scan over: {}.
validation error jax, scan, shape-inference, scalar-operand, pytree
0th dimension of all xs should be replicated. Got {}
validation error jax, scan, sharding, multi-device, jaxmesh
scan got `length` argument of {} which disagrees with leadin
validation error jax, scan, length-mismatch, shape-validation
scan got values with different leading axis sizes: {}.
validation error jax, scan, shape-mismatch, length-inference
scan got no values to scan over and `length` not provided.
validation error jax, scan, missing-length, argument-validation
{} function carry input and carry output must have the same
validation error jax, pytree, scan, while-loop, carry-structure
{} function carry input and carry output must have equal typ
validation error jax, dtype-mismatch, scan, while-loop, shard-map, vma
scan number of arguments doesn't match the number of jaxpr a
validation error jax, internal, scan, jaxpr, argument-count
Scan carry input and output got mismatched varying manual ax
validation error jax, shard-map, vma, scan, spmd, multi-device
lax.while_loop: body_fun and cond_fun arguments should be ca
validation error jax, while-loop, typeerror, callable, api-misuse
cond_fun must return a boolean scalar, but got pytree {}.
validation error jax, while-loop, cond-fun, pytree, boolean-scalar
cond_fun must return a boolean scalar, but got output type(s
validation error jax, while-loop, cond-fun, dtype, scalar
Effects not supported in `while`: {}
error_code error jax, effects, while-loop, debug-print, not-implemented
Ordered IO effects not supported in vmap.
error_code error jax, vmap, ordered-io, debug-print, while-loop, batching
Unordered IO effects not supported in while_loop with batche
error_code error jax, vmap, while-loop, io-effects, control-flow
Reverse-mode differentiation does not work for lax.while_loo
validation error jax, autodiff, grad, while-loop, fori-loop, dynamic-bounds
Cannot write to the same ref in both cond and body of while
error_code error jax, state-api, refs, while-loop, state-discharge
Body jaxpr has consts. If you see this error, please open an
error_code error jax, state-discharge, internal-error, while-loop
Cond jaxpr has consts. If you see this error, please open an
error_code error jax, state-discharge, internal-error, while-loop
lax.fori_loop: body_fun argument should be callable.
validation error jax, fori-loop, typeerror, argument-validation
lower and upper arguments to fori_loop must have equal types
validation error jax, fori-loop, dtype-mismatch, typeerror
Can only use `unroll` in `fori_loop` if the loop bounds are
validation error jax, fori-loop, unroll, static-shape, traced-values
0th dimension of leaf passed to `jax.lax.map` should be repl
validation error jax, lax-map, sharding, multi-device, gspmd
lax.associative_scan: fn argument should be callable.
validation error jax, associative-scan, type-error, callable
associative scan over axis of non-constant size: {}. You may
error_code error jax, associative-scan, dynamic-shape, shape-polymorphism
Array inputs to associative_scan must have the same first di
validation error jax, associative-scan, shape-mismatch
XLA operations do not allow negative axes
validation error jax, lax, cumsum, negative-axis
axis {} is out of bounds for array of shape {}
validation error jax, lax, cumsum, axis-out-of-bounds
{} does not accept dtype {}. Accepted dtypes are subtypes of
validation error jax, lax, dtype, cumsum
{}() output shapes must match {}, got {} and {}
validation error jax, custom-linear-solve, shape-mismatch, linear-algebra
transpose_solve required for backwards mode automatic differ
validation error jax, autodiff, custom-linear-solve, transpose
open an issue at https://github.com/google/jax !!
error_code error jax, autodiff, custom-linear-solve, not-implemented
String padding is not implemented for transposed convolution
validation error jax, lax, convolution, transposed-conv, padding
padding argument to conv_general_dilated should be a string
validation error jax, lax, convolution, padding
Invalid padding mode: {padding}
validation error jax, lax, conv-transpose, padding
No 4+ dimensional dimension_number defaults.
validation error jax, lax, conv-transpose, dimension-numbers, high-rank
conv_general_dilated lhs and rhs must have the same number o
validation error jax, lax, convolution, shape-mismatch
conv_general_dilated feature_group_count must be a positive
validation error jax, lax, convolution, feature-group-count
conv_general_dilated feature_group_count must divide lhs fea
validation error jax, lax, convolution, feature-group-count
conv_general_dilated lhs feature dimension size divided by f
validation error jax, lax, convolution, feature-group-count, kernel-shape
conv_general_dilated rhs output feature dimension size must
validation error jax, lax, convolution, feature-group-count, kernel-shape
conv_general_dilated batch_group_count must be a positive in
validation error jax, lax, convolution, batch-group-count
conv_general_dilated batch_group_count must divide lhs batch
validation error jax, convolution, shape-validation, grouped-conv
conv_general_dilated rhs output feature dimension size must
validation error jax, convolution, shape-validation, grouped-conv
At most one of batch_group_count and feature_group_count may
validation error jax, convolution, grouped-conv, api-misuse
conv_general_dilated window and window_strides must have the
validation error jax, convolution, strides, shape-validation
Convolutions with non-static strides, dilation, feature_grou
exception error jax, jit, static-shapes, convolution, lowering
Wrong number of explicit pads for convolution: expected {},
validation error jax, convolution, padding, shape-validation
Negative padding is larger than the size of the correspondin
validation error jax, convolution, padding, shape-validation
convolution requires lhs and rhs ndim to be equal, got {} an
validation error jax, convolution, shape-validation, rank-mismatch
convolution dimension_numbers list/tuple must be length 3, g
validation error jax, convolution, dimension-numbers, api-misuse
convolution dimension_numbers elements must be strings, got
validation error jax, convolution, dimension-numbers, type-error
convolution dimension_numbers[{}] must have len equal to the
validation error jax, convolution, dimension-numbers, shape-validation
convolution dimension_numbers must be tuple/list or None, go
validation error jax, convolution, dimension-numbers, type-error
convolution dimension_numbers[{}] must contain the character
validation error jax, convolution, dimension-numbers, validation
convolution dimension_numbers[{}] cannot have duplicate char
validation error jax, convolution, dimension-numbers, validation
convolution dimension_numbers elements must each have the sa
validation error jax, convolution, dimension-numbers, validation
Unknown FFT type '{s}'
validation error jax, fft, api-misuse, invalid-enum-value
Unknown FFT type value '{fft_type}'
validation error jax, fft, type-error
only real valued inputs supported for rfft
validation error jax, fft, dtype, api-misuse
full must be called with scalar fill_value, got fill_value.s
validation error jax, lax, type-error, scalar-required, shape-validation
offset must be an integer, got {offset!r}
validation error jax, lax, tri, integer-required, type-error
repeated axis in lax.expand_dims: {dimensions}
validation error jax, lax, expand-dims, duplicate-axis, value-error
repeated axis in lax.expand_dims: {dims}
validation error jax, lax, expand-dims, axis-canonicalization, duplicate-axis
Invalid dimension range passed to collapse: {operand.shape}[
validation error jax, lax, collapse, dimension-range, value-error
Arguments to batch_matmul must be at least 2D, got {}, {}
validation error jax, lax, batch-matmul, dimensionality, value-error
Arguments to batch_matmul must have same ndim, got {}, {}
validation error jax, lax, batch-matmul, rank-mismatch, value-error
Called {name} with a float0 array. float0s do not support an
validation error jax, float0, autodiff, gradient, type-error
{} does not accept dtype {}. Accepted dtypes are subtypes of
exception error jax, lax, dtype-validation, type-error
{name} does not accept dtype {dtype_to_string(aval.dtype)}.
exception error jax, lax, narrow-int, int4, dtype-validation
unreduced rule for {name} is not implemented. Please file an
exception error jax, sharding, gspmd, not-implemented, unreduced
Called {name} with a float0 at position {i}. float0s do not
exception error jax, float0, autodiff, binop, type-error
{} does not accept dtype {} at position {}. Accepted dtypes
exception error jax, lax, binop, dtype-validation, type-error
{}: arrays must have same number of dimensions, got {}.
exception error jax, sharding, named-sharding, rank-mismatch, type-error
transpose with implicit broadcasting of unshaped values. Got
exception error jax, autodiff, vjp, dynamic-shapes, transpose, type-error
the first argument to pow must have an inexact dtype (float
exception error jax, lax, pow, dtype-mismatch, type-error
Integers cannot be raised to negative powers, got integer_po
exception error jax, lax, integer-pow, negative-exponent, type-error
The denominator cannot be unreduced passed to `div`. Got {y=
exception error jax, sharding, div, unreduced, value-error
Denominator should be reduced along the same axes numerator
exception error jax, sharding, lax, division, distributed
HLO comparison {direction} for extended dtype {avals_in[0].d
exception error jax, extended-dtype, comparison, lax, hlo
`sharding` passed to `convert_element_type` can only contain
exception error jax, sharding, dtype-conversion, lax
can only convert to an extended dtype on an array type,but g
exception error jax, extended-dtype, abstract-eval, dynamic-shape, typeerror
Cannot convert_element_type from {dtype_to_string(x.dtype)}
exception error jax, extended-dtype, dtype-conversion, validation
can only convert to extended dtype from its representation d
exception error jax, extended-dtype, dtype-conversion, representation
can only convert to extended dtype from an array of its repr
exception error jax, extended-dtype, shape-validation, rank-mismatch
can only convert to extended dtype from an array of its repr
exception error jax, extended-dtype, shape-validation, suffix-mismatch
can only convert to extended dtype from an array with traili
exception error jax, extended-dtype, sharding, partition-spec
can only convert from an extended dtype on an array type,but
exception error jax, extended-dtype, abstract-eval, typeerror
Cannot convert_element_type from {dtype_to_string(x.dtype)}
exception error jax, extended-dtype, dtype-conversion, validation
can only convert from extended dtype to its representation d
exception error jax, extended-dtype, dtype-conversion, representation
Attempting to convert array of shape {operand.shape} from {o
exception error jax, bitcast, dtype, shape-validation
lax.bitcast_convert_type does not support bool or complex va
exception error jax, bitcast, complex, bool, typeerror
Input type is incompatible with `preferred_element_type`. Th
exception error jax, dot-general, preferred-element-type, dtype-validation
`preferred_element_type` must have the same signedness as th
exception error jax, dot-general, signedness, dtype-validation
`preferred_element_type` must not be narrower than the origi
exception error jax, dot-general, precision, dtype-validation
dot_general requires lhs dimension numbers to be nonnegative
exception error jax, dot-general, dimension-numbers, index-out-of-range
dot_general requires rhs dimension numbers to be nonnegative
exception error jax, dot-general, dimension-numbers, index-out-of-range
dot_general requires equal numbers of lhs_batch and rhs_batc
exception error jax, dot-general, batch-dims, dimension-numbers
dot_general requires lhs batch dimensions to be distinct, go
exception error jax, dot-general, dimension-numbers, batch-dims
dot_general requires rhs batch dimensions to be distinct, go
exception error jax, dot-general, dimension-numbers, batch-dims
dot_general requires lhs contracting dimensions to be distin
exception error jax, dot-general, dimension-numbers, contraction
dot_general requires rhs contracting dimensions to be distin
exception error jax, dot-general, dimension-numbers, contraction
dot_general requires lhs batch dimensions to be disjoint fro
exception error jax, dot-general, dimension-numbers, batch-dims, contraction
dot_general requires rhs batch dimensions to be disjoint fro
validation error jax, dot-general, dimension-numbers, batch-dims, contraction
dot_general requires lhs batch dimensions and rhs batch dime
validation error jax, dot-general, shape-mismatch, batch-dims
dot_general requires contracting dimensions to have the same
validation error jax, dot-general, shape-mismatch, contraction
`out_sharding` passed to `dot_general` can only contain unre
validation error jax, dot-general, sharding, spmd, experimental-api
{name} argument type error: {lhs.dtype}, {rhs.dtype}
validation error jax, dot-general, dtype-mismatch, mixed-precision
`preferred_element_type` must not be narrower than the origi
validation error jax, dot-general, preferred-element-type, dtype-narrowing
Accuracy {accuracy} not supported
exception error jax, dot-general, accuracy, experimental-api
The precision '{precision}' is not supported by dot_general
validation error jax, dot-general, cpu-backend, precision, fp8
lhs_ragged_dim {lhs_ragged_dim} not found in lhs_noncontract
validation error jax, ragged-dot-general, dimension-numbers
ragged_dot_general requires {dim_name} numbers to be nonnega
validation error jax, ragged-dot-general, index-out-of-range
ragged_dot_general expects exactly one lhs ragged dimension.
validation error jax, ragged-dot-general, dimension-numbers
expected rank of group_sizes to be >=1.
validation error jax, ragged-dot-general, group-sizes, shape-mismatch
expected group_sizes to have shape {expected_gs_shape}, got
validation error jax, ragged-dot-general, group-sizes, shape-mismatch
ragged_dot_general requires the group count (last dimension
validation error jax, ragged-dot-general, dynamic-shapes, jit
ragged_dot_general requires zero group dimensions in the rhs
validation error jax, ragged-dot, shape-validation, moe
ragged_dot_general requires exactly one rhs group dimension
validation error jax, ragged-dot, shape-validation, moe
ragged_dot_general requires rhs group dimension numbers to b
validation error jax, ragged-dot, dimension-numbers, shape-validation
expected rhs group dimension size to be {num_groups}, got {r
validation error jax, ragged-dot, shape-mismatch, moe
ragged_dot_general requires that group_sizes.dtype is subtyp
validation error jax, ragged-dot, dtype-validation
Unimplemented group_offset support.
exception error jax, ragged-dot, autodiff, not-implemented
ragged_dot vmap over any dim but 0 - NYI
exception error jax, ragged-dot, vmap, batching, not-implemented
Explicit sharding inference for ragged_dot_general is not cu
exception error jax, ragged-dot, sharding, multi-device, spmd
broadcast_in_dim broadcast_dimensions must have length equal
validation error jax, broadcast-in-dim, shape-validation
broadcast_in_dim target broadcast shape must have equal or h
validation error jax, broadcast-in-dim, rank-validation
broadcast_in_dim broadcast_dimensions must be a subset of ou
validation error jax, broadcast-in-dim, index-out-of-range
broadcast_in_dim operand dimension sizes must either be 1, o
validation error jax, broadcast-in-dim, shape-mismatch
broadcast_in_dim broadcast_dimensions must not contain dupli
validation error jax, broadcast-in-dim, duplicate-index
`out_sharding` passed to `broadcast_in_dim` can only contain
validation error jax, broadcast-in-dim, sharding, gspmd, unreduced
reps length must be equal to the ndim of x, got {len(reps)=}
validation error jax, tile, shape-validation
clamp requires min.shape == operand.shape or min.shape == ()
validation error jax, clamp, shape-mismatch, broadcasting
clamp requires max.shape == operand.shape or max.shape == ()
validation error jax, clamp, shape-mismatch, broadcasting
concatenate expects at least one operand, got 0.
validation error jax, concatenate, shape-validation, empty-input
All objects to concatenate must be arrays, got {}.
validation error jax, concatenate, type-validation, dtype
Cannot concatenate arrays with different numbers of dimensio
validation error jax, concatenate, rank-mismatch, shape-validation
concatenate dimension out of bounds: dimension {} for shapes
validation error jax, concatenate, axis-out-of-bounds, shape-validation
Cannot concatenate arrays with shapes that differ in dimensi
validation error jax, concatenate, shape-mismatch, shape-validation
stack expects at least one operand, got 0.
validation error jax, stack, empty-input, shape-validation
Cannot stack arrays with different numbers of dimensions: go
validation error jax, stack, rank-mismatch, shape-validation
All input arrays must have the same shape. Got {}.
validation error jax, stack, shape-mismatch
unstack requires arrays with rank > 0, however a scalar arra
validation error jax, unstack, scalar, rank-error
Sizes passed to split must be nonnegative, got {list(sizes)}
validation error jax, split, negative-size, validation
Sum of sizes {np.sum(sizes)} must be equal to dimension {axi
validation error jax, split, size-mismatch, shape-validation
pad operand and padding_value must be same dtype: got {} and
validation error jax, pad, dtype-mismatch
padding_value must be a scalar; got {np.shape(padding_value)
validation error jax, pad, scalar-required, shape-validation
length of padding_config must equal the number of axes of op
validation error jax, pad, config-mismatch, rank-error
interior padding in padding_config must be nonnegative, got
validation error jax, pad, negative-padding, validation
Dimension size after padding is not at least 0, got result s
validation error jax, pad, negative-dimension, shape-validation
dimensions are not unique: {dimensions}
validation error jax, squeeze, duplicate-dims, validation
dimensions outside range [0, ndim): {dimensions}
validation error jax, squeeze, axis-out-of-bounds
cannot select an axis to squeeze out which has size not equa
validation error jax, squeeze, size-one-required, shape-validation
reshape new_sizes must all be positive, got {}.
validation error jax, reshape, negative-dimension, shape-validation
duplicate value in 'axes' of reduction: {axes}
validation error jax, lax, reduction, axes, duplicate
reduction axes {axes} contains out-of-bounds indices for {op
validation error jax, lax, reduction, out-of-bounds, axis
out_sharding passed to {name} cannot be reduced. Got {out_sh
validation error jax, sharding, named-sharding, reduction, distributed
reduce_sum on operand {operand.str_short(True)} is not allow
validation error jax, float16, bfloat16, reduction, mixed-precision, config-flag
Invalid axis {axis} for operand shape {operand.shape}
validation error jax, argmax, argmin, axis, validation
argmin and argmax require non-empty reduced dimension. opera
validation error jax, argmax, empty-array, shape-validation
index_dtype must be an integer type, but got {}
validation error jax, argmin, dtype, type-error
logical reduction requires operand dtype bool or int, got {o
validation error jax, lax, logical-reduction, dtype
reduce_precision: exponent_bits must be positive; got {expon
validation error jax, reduce-precision, bit-manipulation, validation
reduce_precision: mantissa_bits must be non-negative; got {m
validation error jax, reduce-precision, quantization, validation
Arguments to sort must have equal shapes, got: {shapes}
validation error jax, lax, sort, shape-mismatch
top_k is not compatible with complex inputs.
validation error jax, top-k, complex-numbers, dtype
top_k operand must have >= 1 dimension, got {}
validation error jax, top-k, scalar, rank-error
axis argument out of range: {axis=} for {operand.shape=}
validation error jax, top-k, axis, out-of-bounds
k argument to top_k must be no larger than size along axis;
validation error jax, top-k, shape-validation, off-by-one
top_k returns int32 indices, which will overflow for array d
validation error jax, top-k, int32-overflow, large-arrays
Arguments to after_all must be tokens
validation error jax, lax, tokens, control-flow, xla
Arguments to rng_uniform must have identical dtypes, got {}
validation error jax, rng-uniform, dtype-mismatch
Arguments to rng_uniform must be scalars; got shapes {} and
validation error jax, rng-uniform, scalar-requirement, shape-validation
Unsupported dtype: {dtype}
validation error jax, linalg, eig, dtype, gpu
Nonsymmetric eigendecomposition requires cusolver 11.7.1 or
error_code error jax, cuda, cusolver, eig, version-mismatch, gpu
Left eigenvectors are not supported by cusolver
validation error jax, eig, cusolver, gpu, not-implemented
Derivatives of non-symmetric eigenvectors are only valid und
validation error jax, autodiff, eig, gradients, numerical
Argument to symmetric eigendecomposition must have shape [..
validation error jax, linalg, eigh, shape-validation, square-matrix
subset_by_index not supported on CPU and GPU
validation error jax, eigh, subset-by-index, cpu, gpu, not-implemented
QDWH implementation is only supported on TPU
validation error jax, eigh, qdwh, tpu, backend, not-implemented
Jacobi implementation is not supported on CPU
validation error jax, eigh, jacobi, cpu, backend, not-implemented
Derivatives not defined for partial eigen decomposition.
validation error jax, autodiff, eigh, subset-by-index, gradients
Argument to Hessenberg reduction must have shape [..., n, n]
validation error jax, linalg, hessenberg, shape-validation
hessenberg requires the last dimension of a to be constant,
validation error jax, hessenberg, dynamic-shapes, jit, cpu
The first argument to householder_product must have at least
validation error jax, linalg, householder, shape-validation, qr
The second argument to householder_product must not have mor
validation error jax, linalg, householder, shape-validation
ormqr with left=True expects c to have the same number of ro
validation error jax, linalg, ormqr, shape-validation, qr
ormqr with left=False expects c to have the same number of c
validation error jax, linalg, ormqr, shape-validation, qr
'trans' value must be 0, 1, or 2, got {trans}
validation error jax, lu-solve, invalid-argument, api-misuse
last two dimensions of LU decomposition must be equal, got s
validation error jax, lu-solve, shape-validation, square-matrix
b matrix must have rank >= 1, got shape {b.shape}
validation error jax, lu-solve, shape-validation, rank
When LU decomposition matrix and b have the same number of d
validation error jax, lu-solve, shape-validation, broadcasting
When LU decomposition matrix and b different numbers of dime
validation error jax, lu-solve, shape-validation, broadcasting
logaddexp2 requires floating-point or complex inputs; got {x
validation error jax, dtype, lax, validation
This gmm kernel only supports either (m, k) x (g, k, n) -> (
validation error jax, pallas, gpu, shape, moe
Group sizes {group_sizes.shape=} must match first dimension
exception error jax, pallas, moe, shape, gpu
Explicit sharding is not currently supported in the pallas-t
exception error jax, sharding, pallas, spmd, not-implemented
group_offset is not currently supported in the pallas-triton
exception error jax, pallas, gpu, not-implemented, ragged-dot
{name} cannot accept args with unreduced_kind={mat.unreduced
exception error jax, shard-map, spmd, collectives, internal-api
axis_index_groups only supported for sums over just named ax
exception error jax, psum, axis, validation
pargmin only accepts a single axis, got {axis_name}
exception error jax, pargmin, axis, typeerror
axis_index_groups must cover all indices exactly once
exception error jax, psum, validation, axis-index-groups
`perm` passed to `jax.lax.ppermute` must be a list or a tupl
exception error jax, ppermute, typeerror, collectives
`perm` does not represent a permutation: {perm}
exception error jax, pshuffle, permutation, validation
The size of all_to_all split_axis ({x.shape[split_axis]}) ha
exception error jax, all-to-all, shape, spmd
all_to_all requires the size of the mapped axis axis_name to
exception error jax, all-to-all, shape, spmd
axis_index_groups not supported in vmap collectives. Please
exception error jax, vmap, collectives, not-implemented
axis_index_groups can only be used with reductions over name
exception error jax, collectives, axis-validation
Found an unbound axis name: {name}. To fix this, please call
exception error jax, mesh, shard-map, unbound-axis
axis_index_groups must all be the same size for TPU lowering
exception error jax, tpu, collectives, lowering
{op_name} sources and destinations must be unique, got {}.
exception error jax, ppermute, collectives, validation
Make sure that the axis_name passed to jax.lax.ppermute is i
exception error jax, ppermute, mesh, ordering
psend is currently only implemented on GPUs
exception error jax, psend, gpu-only, backend
psend currently only supports manual sharding
exception error jax, psend, shard-map, manual-sharding
psend is currently only implemented on GPU
exception error jax, psend, backend, not-implemented
precv currently only supports manual sharding
exception error jax, precv, shard-map, manual-sharding
precv is currently only implemented on GPU
exception error jax, precv, backend, not-implemented
pbroadcast batcher only supports a single axis
exception error jax, pbroadcast, vmap, not-implemented
Replica groups must be equally sized
exception error jax, all-to-all, collectives, validation
Please open a feature request!
exception error jax, all-to-all, vmap, not-implemented
all_to_all must be used within a mapped context like vmap or
exception error jax, all-to-all, mapped-context
ragged_all_to_all input_offsets must be integer type.
exception error jax, ragged-all-to-all, dtype
ragged_all_to_all send_sizes must be integer type.
exception error jax, ragged-all-to-all, dtype
ragged_all_to_all output_offsets must be integer type.
exception error jax, ragged-all-to-all, dtype
ragged_all_to_all recv_sizes must be integer type.
exception error jax, ragged-all-to-all, dtype
pvary is a invariant->variant collective. This means that th
validation error jax, collectives, mesh, spmd, pvary
all_gather_reduced only accepts inputs that are varying. Got
validation error jax, collectives, mesh, spmd, all-gather
all_gather_reduced is a Varying -> Reduced collective. This
validation error jax, collectives, mesh, axis-name, all-gather
all_gather_reduced's input cannot be reduced across the axis
validation error jax, collectives, mesh, reduced, all-gather
Please file an issue at https://github.com/jax-ml/jax/issues
exception error jax, vmap, collectives, not-implemented, all-gather
unreduced_psum_scatter only accepts inputs that are unreduce
validation error jax, collectives, mesh, spmd, reduce-scatter
unreduced_psum_scatter is a Unreduced -> Varying collective.
validation error jax, collectives, axis-name, mesh, reduce-scatter
unreduced_psum_scatter's input cannot be varying across the
validation error jax, collectives, mesh, varying, reduce-scatter
{name} only accepts inputs that are unreduced. Got {aval.str
exception error jax, sharding, collectives, validation, mesh
{name} is a Unreduced -> Invariant collective. This means th
exception error jax, mesh, axis-name, collectives, sharding
{name}'s input cannot be varying across the axis_name provi
exception error jax, sharding, collectives, varying, mesh
{name} does not accept integer axis_name. Got axis_name={axe
exception error jax, axis-name, collectives, api-misuse
preduced is a Invariant->Reduced collective. This means that
exception error jax, sharding, collectives, transpose, mesh
preduced input cannot be reduced across the axis_name provid
exception error jax, sharding, collectives, double-reduction, mesh
vary_unreduced_cast only accepts inputs that are varying. Go
exception error jax, named-axes, manual-parallelism, aval-validation
vary_unreduced_cast is a Varying->Unreduced collective. This
exception error jax, axis-name, named-axes
vary_unreduced_cast input cannot be unreduced across the axi
exception error jax, named-axes, duplicate-cast
reduced_vary_cast only accepts inputs that are reduced. Got
exception error jax, named-axes, manual-parallelism, aval-validation
reduced_vary_cast is a Reduced->Varying collective. This mea
exception error jax, axis-name, named-axes
reduced_vary_cast input cannot be varying across the axis_na
exception error jax, named-axes, manual-parallelism
{name} can only accept axis_name which corresponds to one of
exception error jax, pcast, named-axes, mixed-state
{axis_name=} must be a tuple or a str. Got {axis_name}
exception error jax, pcast, type-error, axis-name
Got unexpected `to` value. Allowed `to` values are: {_allowe
exception error jax, pcast, invalid-argument
Unsupported pcast from={from_}, {to=}
exception error jax, pcast, unsupported-transition, named-axes
Unsupported type for JSON serialization: {type(obj)} ({obj})
exception error jax, json-serialization, async-collectives, config
async done op got {aval}, want core.AbstractFuture
exception error jax, async-collectives, future, type-error
{side} contracting dim {i} of size {size} must be divisible
exception error jax, scaled-dot, shape-validation, quantization
The ratio of {side} contracting dim {i} to its scale's dim s
exception error jax, scaled-dot, shape-validation
{side} dim {i} of size {size} does not match scale dim size
exception error jax, scaled-dot, shape-validation
All input tensors must have the same rank. Got lhs rank: {lh
exception error jax, scaled-dot, rank-validation
LHS and RHS must have the same number of batch dimensions, g
exception error jax, scaled-dot, dimension-numbers
dynamic_update_slice update shape must be smaller than opera
exception error jax, lax, dynamic-update-slice, shape-mismatch
start_indices arguments to dynamic_update_slice must be scal
exception error jax, lax, dynamic-update-slice, scalar-indices
index arguments to dynamic_update_slice must be integers of
exception error jax, lax, dtype, indices
indices must have an integer type
exception error jax, lax, gather, dtype
{name} in {op_name} op must be sorted; got {dims}
exception error jax, lax, gather, scatter, dimension-numbers
Invalid {name} set in {op_name} op; valid range is [0, {rank
exception error jax, lax, gather, scatter, dimension-out-of-range
Invalid {name} set in {op_name} op; valid range is [0, {rank
exception error jax, lax, gather, scatter, dimension-out-of-range
{name} in {op_name} op must not repeat; got: {dims}.
exception error jax, lax, gather, scatter, duplicate-dims
{name1} and {name2} in {op_name} op must be disjoint; got: {
exception error jax, lax, gather, scatter, dimension-numbers
Gather index leaf dimension must be within [0, rank(indices)
exception error jax, lax, gather, index-vector-dim
reduce_window expected init_values to be scalars but init_va
validation error jax, reduce-window, shape-validation, scalar
reduce_window batching is not implemented for initial values
exception error jax, vmap, batching, not-implemented, reduce-window
reduce_window jvp does not support non-zero init_value_tange
validation error jax, autodiff, jvp, reduce-window
Cannot lower effectful `reduce_window`.
exception error jax, effects, lowering, mlir, reduce-window
operand to reduce_window_sum must have a number dtype, got {
validation error jax, dtype, reduce-window, type-validation
reduce_window got the wrong number of window_dimensions for
validation error jax, shape-validation, reduce-window, windowing
reduce_window got inconsistent window_strides and window_dim
validation error jax, shape-validation, reduce-window, strides
reduce_window got inconsistent base_dilation and window_dime
validation error jax, shape-validation, reduce-window, dilation
reduce_window got inconsistent window_dilation and window_di
validation error jax, shape-validation, reduce-window, dilation
select_and_scatter got inconsistent window_strides and windo
validation error jax, shape-validation, select-and-scatter, pooling
select_and_gather_add tangents and operand shapes must match
validation error jax, autodiff, vjp, pooling, shape-validation
VJP not implemented for select_and_gather (MaxPool) with win
error_code error jax, autodiff, vjp, pooling, dilation, not-implemented
mulhi requires integer inputs, got {dtype}
validation error jax, dtype, integer, mulhi, type-validation
mulhi operands must have the same dtype, got {dtype} and {y_
validation error jax, dtype, mulhi, type-mismatch
Dimension mismatch
validation error jax, numpy-reference, convolution, shape-validation
Wrong number of strides for spatial dimensions
validation error jax, numpy-reference, strides, shape-validation, convolution
Wrong number of pads for spatial dimensions
validation error jax, numpy-reference, padding, shape-validation, convolution
Expected mode of type `LayoutMode`. Got type: {type(mode)}
validation error jax, layout, type-validation, enum
Length of major_to_minor and the rank of the value should ma
validation error jax, layout, rank-mismatch, validation
Sharding has to be concrete when layout is of type {type(lay
validation error jax, layout, sharding, validation
Invalid value received for the layout argument. Expected val
validation error jax, layout, type-validation, isinstance-check
Invalid value received for the sharding argument. Expected v
validation error jax, sharding, type-validation
Cannot determine the ``__name__`` of the caller.
error_code error jax, lazy-loading, import-machinery, internal-api
module '{package_name}' has no attribute '{name}'
validation error jax, lazy-loading, attributeerror, module-structure
Attempting to set log level "{logging_level}" which isn't on
validation error jax, logging, config, valueerror
Please install the `etils[epath]` package to specify a cache
error_code error jax, compilation-cache, optional-dependency, cloud-storage
Please install the `filelock` package to set `jax_compilatio
error_code error jax, compilation-cache, optional-dependency, filelock
key cannot be empty
validation error jax, lru-cache, input-validation
Cannot update the mesh of the current resource environment.
validation error jax, mesh, axis-names, context-manager
When passing host local inputs to pjit, devices connected to
validation critical jax, mesh, multi-host, tpu, device-topology
axis_types passed to {name} must be of type `jax.sharding.Ax
validation error jax, mesh, axis-types, type-validation
Number of axis names should match the number of axis_types.
validation error jax, mesh, axis-types, shape-mismatch
Mesh axis names cannot be None. Got: {axis_names}
validation error jax, mesh, axis-names, input-validation
Mesh requires the ndim of its first argument (`devices`) to
validation error jax, mesh, shape-mismatch, input-validation
Mesh context manager is disabled.
error_code error jax, mesh, deprecation, context-manager, config-flag
AbstractMesh does not implement {name}
validation error jax, mesh, abstract-class, not-implemented
Expected mesh of type `jax.sharding.AbstractMesh`. Got type:
validation error jax, mesh, type-validation, abstract-mesh
use_abstract_mesh cannot change the size of the mesh. Got ne
validation error jax, mesh, abstract-mesh, size-mismatch
Failed to find assignment for logical_axis_index {logical_ax
exception error jax, mesh, tpu, topology, not-implemented
Reduced axes can only refer to mesh axes that is of type `Ex
validation error jax, sharding, mesh, distributed, named-sharding
Unknown algorithm '{algorithm}'. Expected 'fast' or 'stable'
validation error jax, nn, standardize, enum-argument
Expected num_classes to match the size of axis {axis}, but {
validation error jax, nn, one-hot, pmap, distributed
The input part of spec in out_sharding should match the spec
validation error jax, nn, one-hot, sharding, named-sharding
{name} ndim should be {len(shape)}, but got {t.ndim}
validation error jax, nn, attention, shape-validation
{name} dtype should be {dtype}, but got {t.dtype}
validation error jax, nn, attention, dtype-validation
{name} shape should be {shape}: but got {t.shape}
validation error jax, nn, attention, shape-validation
The number of query heads must be a multiple of key/value he
validation error jax, nn, attention, gqa, shape-validation
cuDNN doesn't support right window: {r_window} when causal m
validation error jax, nn, attention, cudnn, sliding-window, gpu
Unsupported implementation option: {implementation}
validation error jax, nn, attention, enum-argument
scaled_matmul requires all inputs to be 3-dimensional arrays
validation error jax, nn, matmul, float8, mxfp8, shape-validation
scaled_matmul requires inputs a and b to have matching batch
validation error jax, nn, matmul, float8, shape-validation
scaled_matmul requires scales to have matching batch (B) and
validation error jax, nn, matmul, float8, scales, shape-validation
scaled_matmul requires scales to match non-contract dimensio
validation error jax, nn, matmul, float8, scales, shape-validation
Unsupported mode: {mode}
validation error jax, nn, float8, config, enum-argument
Can't compute input and output sizes of a {len(shape)}-dimen
validation error jax, nn, initializers, shape-validation
invalid mode for variance scaling initializer: {mode}
exception error jax, nn, initializers, enum-argument
invalid distribution for variance scaling initializer: {dist
exception error jax, nn, initializers, enum-argument
orthogonal initializer requires at least a 2D shape
exception error jax, nn, initializers, shape-validation
Delta orthogonal initializer requires a 3D, 4D or 5D shape.
exception error jax, initializer, shape-validation
`fan_in` must be less or equal than `fan_out`.
exception error jax, initializer, shape-validation, fan-in-out
{api_version=!r} is not available; available versions are: {
exception error jax, array-api, version-mismatch
Currently, string arrays can only be made from NumPy arrays.
exception error jax, string-dtype, type-validation
Cannot make an array with dtype {dtype} from an object with
exception error jax, string-dtype, dtype-mismatch
ndmin {ndmin} cannot be greater than object's ndims {object.
exception error jax, string-dtype, shape-validation
array() takes at most 5 positional arguments but {len(args)
exception error jax, argument-validation, numpy-compat
array() got multiple values for argument '{name}'
exception error jax, argument-validation, numpy-compat, deprecation
Only implemented for order='K'
exception error jax, not-implemented, numpy-compat, memory-order
None is not a valid value for jnp.array
exception error jax, null-handling, data-validation
Unexpected input type for array: {type(object)}
exception error jax, type-validation, conversion
`{device_or_sharding = }` was passed to`canonicalize_or_get_
exception error jax, device-validation, type-validation
jnp.asarray: cannot convert object of type {type(a)} to JAX
exception error jax, copy-semantics, device-placement, array-api
expected sequence object with len >= 0 or a single integer
exception error jax, shape-validation, generator
Cannot interpret '{dtype}' as a data type.\n\nDid you accide
exception error jax, shape-validation, dtype, typo-detection
Cannot interpret '{dtype}' as a data type.\n\nDid you accide
exception error jax, shape-validation, dtype, typo-detection
Cannot interpret '{dtype}' as a data type.\n\nDid you accide
exception error jax, shape-validation, dtype, typo-detection
Number of samples, {num}, must be non-negative.
exception error jax, linspace, argument-validation
Array.__contains__: unsupported operand type {type(other)}.
exception error jax, membership, typeerror, contains
Array.__contains__: search array must be one-dimensional, go
exception error jax, membership, shape-mismatch, contains
Array.__contains__: query value must be a scalar, got {query
exception error jax, membership, scalar-required, contains
No Python scalar type for {arr.dtype=}
exception error jax, dtype, item, extended-dtype
np.reshape order=A is not implemented.
exception error jax, reshape, not-implemented, numpy-compat
Unexpected value for 'order' argument: {order}.
exception error jax, reshape, invalid-argument, validation
stream argument of array.to_device()
exception error jax, device-transfer, stream, not-implemented
can only specify one unknown axis size with a `-1` value, go
exception error jax, reshape, invalid-shape
cannot reshape array of shape {arr.shape} (size {arr.size})
exception error jax, reshape, size-mismatch
cannot reshape array of shape {arr.shape} (size {arr.size})
exception error jax, reshape, size-mismatch
`type` argument of array.view() is not supported.
exception error jax, view, not-implemented, numpy-compat
When changing to a larger dtype, its size must be a divisor
exception error jax, view, dtype, alignment
JAX Arrays do not implement the arr.flat property: consider
exception error jax, flat, not-implemented, numpy-compat
dtypes of sorted_arr and query must match; got {sorted_arr_a
exception error jax, dtype, searchsorted, validation
invalid argument side={side!r}, expected 'left' or 'right'
exception error jax, searchsorted, argument-validation
invalid argument {method=}, expected one of {list(self.valid
exception error jax, searchsorted, argument-validation
batch dimension sizes must match; got {sorted_arr_aval.shape
exception error jax, shape-mismatch, searchsorted, batching
out_dtype should be an integer type; got {out_dtype}
exception error jax, dtype, searchsorted, argument-validation
size must be a positive integer; got {size=}
exception error jax, nonzero, argument-validation
out_dtype must be integer typed; got {out_dtype=}
exception error jax, dtype, nonzero, argument-validation
axes out of range for array with {a_aval.ndim} dimensions:
exception error jax, nonzero, axes, argument-validation
duplicate axes are not allowed: {axes=}
exception error jax, nonzero, axes, argument-validation
Expected {len(axes)} fill values, got {len(fill_value_avals)
exception error jax, nonzero, fill-value, argument-validation
Expected fill values to have dtype {out_dtype}, got {fill_va
exception error jax, nonzero, fill-value, dtype
fill_value shape {fv_aval.shape} is not broadcast-compatible
exception error jax, nonzero, fill-value, broadcasting
fill_value shape {fv_aval.shape} cannot be broadcast to batc
exception error jax, nonzero, fill-value, broadcasting
Run out of characters for batch dimension in einsum.
exception error jax, einsum, vmap, resource-limit
Transpose of Einsum with multiple linear inputs is not suppo
exception error jax, einsum, autodiff, not-implemented
Unsupported method: {method}
exception error jax, searchsorted, internal-api, argument-validation
fill_value tuple must have length equal to number of axes ({
validation error jax, nonzero, fill-value, argument-validation
fill_value must be a scalar or tuple of scalars; got {fill_v
validation error jax, nonzero, fill-value, argument-validation
could not understand directive {directive!r}
validation error jax, r-c-index-tricks, directive, numpy-compat
string directive must be placed at the beginning
validation error jax, r-c-index-tricks, directive, numpy-compat
Indexer must have integer or boolean type, got indexer with
validation error jax, indexing, dtype
JAX does not support string indexing; got {idx=}
validation error jax, indexing, string-index
Indexer must have integer or boolean type, got indexer with
validation error jax, indexing, dtype, sequence
Indexer must have integer or boolean type, got indexer with
validation error jax, indexing, scalar, dtype
only integers, slices (`:`), ellipsis (`...`), newaxis (`Non
validation error jax, indexing, type-error
Unrecognized index type: {typ}
validation error jax, internal, invariant
an index can only have a single ellipsis ('...')
validation error jax, indexing, ellipsis
Too many indices: array is {len(shape)}-dimensional, but {to
validation error jax, indexing, rank, shape-mismatch
index {i} out of bounds for axis {axis} with size {size} ({n
validation error jax, indexing, out-of-bounds
Array slice indices must have static start/stop/step to be u
validation error jax, jit, tracing, dynamic-slice
boolean index did not match shape of indexed array in index
validation error jax, boolean-mask, shape-mismatch, indexing
mode='slice' is not valid for polymorphic shapes.
validation error jax, polymorphic-shapes, export, slicing
static_slice requires mode='promise_in_bounds' or mode='clip
validation error jax, indexing, mode, api-misuse
static_slice with partial slices does not support nontrivial
validation error jax, sharding, multi-device, slicing
static_slice: indices must be static scalars or slices. Got
validation error jax, static-slice, advanced-indexing
static_slice: unrecognized index {pidx.index} at position {p
validation error jax, internal, invariant
index is out of bounds for axis {axis} with size 0
validation error jax, empty-array, out-of-bounds, indexing
static_slice: unrecognized index {pidx.index}
validation error jax, internal, invariant
dynamic_slice requires mode='promise_in_bounds' or mode='cli
validation error jax, dynamic-slice, mode, api-misuse
dynamic_slice with partial slices does not support nontrivia
validation error jax, sharding, dynamic-slice, multi-device
dynamic_slice: only unit steps supported in slice. Got {pidx
validation error jax, indexing, slice, stride, dynamic-slice
dynamic_slice: only scalar indices allowed. Got index of typ
validation error jax, indexing, scalar-index, dynamic-slice
dynamic_slice: indices must be scalars or slices. Got index
validation error jax, indexing, boolean-index, dynamic-slice
dynamic_slice: unrecognized index {pidx.index} at position {
validation error jax, indexing, unsupported-index, dynamic-slice
dynamic_slice: unrecognized index {pidx.index}
validation error jax, indexing, internal, dynamic-slice
The 'out' argument to jnp.take is not supported.
validation error jax, numpy-compat, out-arg, immutable
The 'raise' mode to jnp.take is not supported.
validation error jax, take, mode, bounds-check
Invalid mode '{mode}' for np.take
validation error jax, take, invalid-argument, mode
Cannot do a non-empty jnp.take() from an empty axis.
validation error jax, take, empty-array, out-of-bounds
take_along_axis indices must be of integer type, got {index_
validation error jax, take-along-axis, dtype, integer-index
take_along_axis indices must be 1D if axis=None, got shape {
validation error jax, take-along-axis, shape-mismatch, axis
indices and arr must have the same number of dimensions; {}
validation error jax, take-along-axis, ndim, shape-mismatch
One Hot indexing is only supported for up to 50 leading dime
exception error jax, take-along-axis, rank-limit, einsum
jax.numpy.put_along_axis cannot modify arrays in-place, beca
exception error jax, put-along-axis, immutability, inplace
put_along_axis arguments 'arr' and 'indices' must have same
exception error jax, put-along-axis, ndim, shape-mismatch
put_along_axis argument 'values' must be broadcastable to 'i
exception error jax, put-along-axis, broadcasting, shape-mismatch
Expected strategy to be IndexingStrategy; got {strategy}
exception error jax, internal, indexing-strategy, type-error
fill_value argument to indexed get() must be a scalar
exception error jax, indexing, fill-value, scalar
the only valid string value of `left` is 'extrapolate', but
exception error jax, interp, invalid-argument, sentinel-value
the only valid string value of `right` is 'extrapolate', but
exception error jax, interp, invalid-argument, sentinel-value
jnp.interp: complex x values not supported.
exception error jax, interp, complex-dtype, unsupported-operation
period must be a scalar; got {period}
exception error jax, interp, period, shape-validation
size and fill_value arguments cannot be used in three-term w
exception error jax, where, invalid-argument, api-misuse
Either both or neither of the x and y arguments should be pr
exception error jax, where, missing-argument
condlist must have length equal to choicelist ({} vs {})
exception error jax, select, length-mismatch
condlist must be non-empty
exception error jax, select, empty-list
x argument to bincount must have an integer type; got {x.dty
exception error jax, bincount, dtype, typeerror
only 1-dimensional input supported.
exception error jax, bincount, ndim, shape-validation
type of weights must match type of x. Got typeof(x)={core.ty
exception error jax, bincount, weights, sharding, shape-mismatch
jax.numpy.{op} with a symbolic number of sections is not sup
exception error jax, split, symbolic-dim, tracing
array split does not result in an equal division: rest is {r
exception error jax, split, uneven-division
No input was provided to the clip function.
exception error jax, clip, missing-argument
Clip received a complex value either through the input or th
exception error jax, clip, complex-dtype, unsupported-operation
The 'out' argument to jnp.round is not supported.
exception error jax, round, out-parameter, notimplementederror
integer np.round not implemented for decimals < 0
exception error jax, round, integer-dtype, notimplementederror
fill_value must be a scalar or a tuple of length {arr.ndim};
exception error jax, nonzero, fill-value, scalar-validation
jnp.unwrap does not support complex inputs.
exception error jax, unwrap, complex-dtype, unsupported-operation
`{name}` entries must be the same shape: {nvals}
exception error jnp-pad, shape-mismatch, numpy
jnp.pad: {name} with {nd=} has unsupported shape {nvals.shap
exception error jnp-pad, invalid-argument-shape
Cannot apply '{}' padding to empty axis
exception error jnp-pad, empty-axis, reflect, wrap
jnp.pad: constant_values has unsupported shape {constant_val
exception error jnp-pad, constant-values, invalid-argument-shape
Shape polymorphism is supported for jnp.pad with 'reflect' o
exception error jnp-pad, shape-polymorphism, symbolic-dimensions, reflect
stat_length of 0 yields no value for padding
exception error jnp-pad, stat-length, statistical-mode
Expected pad_width to have shape {(nd, 2)}; got {pad_width_a
exception error jnp-pad, pad-width, shape-validation
index can't contain negative values
exception error jnp-pad, negative-pad-width
`pad_width` must be of integral type.
exception error jnp-pad, type-error, integral-type
Unimplemented padding mode '{}' for np.pad.
exception error jnp-pad, unsupported-mode, not-implemented
unsupported keyword arguments for mode '{}': {}
exception error jnp-pad, unsupported-kwarg, keyword-arguments
Need at least one array to stack.
validation error jnp-stack, empty-sequence, value-error
The 'out' argument to jnp.stack is not supported.
exception error jnp-stack, out-argument, immutable-arrays, numpy-parity
Need at least one array to concatenate.
validation error jnp-concatenate, empty-sequence, value-error
Zero-dimensional arrays cannot be concatenated.
validation error jnp-concatenate, wrong-argument-form
The 'out' argument to jnp.choose is not supported.
exception error jnp-choose, out-argument, immutable-arrays, numpy-parity
`a` array must be integer typed
validation error jax, numpy, dtype, valueerror, choose
invalid entry in choice array
validation error jax, choose, index-out-of-range, valueerror
mode={mode!r} not understood. Must be 'raise', 'wrap', or 'c
validation error jax, choose, invalid-argument, valueerror
jax.numpy.block does not allow tuples, got {}
validation error jax, block, tuple-vs-list, valueerror
jax.numpy.block does not allow empty list arguments
validation error jax, block, empty-list, valueerror
Mismatched list depths in jax.numpy.block
validation error jax, block, shape-mismatch, valueerror
jnp.fromfile() is not implemented because it may be non-pure
exception error jax, fromfile, notimplemented, io, jit-purity
jnp.fromiter() is not implemented because it may be non-pure
exception error jax, fromiter, notimplemented, io, jit-purity
k must be a scalar integer; got {k}
validation error jax, eye, dtype, scalar, valueerror
negative dimensions are not allowed, got {N} and {M}
validation error jax, eye, negative-dimension, valueerror
jax.numpy.arange: arguments must be scalars; got {name}={val
validation error jax, arange, scalar-requirement, valueerror
Passing complex start/stop/step to jnp.arange is no longer s
validation error jax, arange, complex, breaking-change, v0-10
In arange with non-constant arguments all of start, stop, an
validation error jax, arange, symbolic-dim, shape-polymorphism, valueerror
arange has step == 0
validation error jax, arange, zero-step, valueerror
jax.numpy.meshgrid only supports copy=True
validation error jax, meshgrid, copy-semantics, valueerror
Valid values for indexing are 'xy' and 'ij', got {indexing}
validation error jax, meshgrid, indexing, invalid-argument
Unsupported input type to jax.numpy.i0: {x_arr.dtype}
validation error jax, i0, bessel, dtype, complex
np.delete(arr, obj): for boolean indices, obj must be one-di
exception error jax, numpy-compat, boolean-mask, shape-mismatch
np.delete(arr, obj): got obj.dtype={obj_array.dtype}; must b
exception error jax, dtype, indexing
jnp.insert(): obj must be a slice, a one-dimensional array,
exception error jax, numpy-compat, insert, ndim
jnp.insert(): index array must be integer typed; got {obj}
exception error jax, insert, dtype, indexing
function is not returning an array of the correct shape
exception error jax, apply-along-axis, shape-mismatch
Dimension must be either 2 or 3 for cross product
exception error jax, cross-product, shape-validation
x must be a one-dimensional array
exception error jax, vander, shape-validation
N must be nonnegative
exception error jax, vander, argument-validation
The 'out' argument to jnp.argmax is not supported.
exception error jax, argmax, unsupported-argument, numpy-compat
attempt to get argmax of an empty sequence
exception error jax, argmax, empty-array
The 'out' argument to jnp.argmin is not supported.
exception error jax, argmin, unsupported-argument, numpy-compat
attempt to get argmin of an empty sequence
exception error jax, argmin, empty-array
The 'out' argument to jnp.nanargmax is not supported.
exception error jax, nanargmax, unsupported-argument, numpy-compat
The 'out' argument to jnp.nanargmin is not supported.
exception error jax, nanargmin, unsupported-argument, numpy-compat
'shift' and 'axis' arguments to roll must be scalars or 1D a
exception error jax, roll, shape-validation
{start=} must satisfy {-a_ndim}<=start<={a_ndim}
exception error jax, rollaxis, argument-validation
Expected an input array of integer or boolean data type
exception error jax, packbits, dtype-validation
'order' must be either 'little' or 'big'
exception error jax, packbits, bitorder, argument-validation
Expected an input array of unsigned byte data type
exception error jax, unpackbits, dtype-validation
Arguments to jax.numpy.gcd must be integers.
exception error jax, numpy, dtype-validation, integer-required
Arguments to jax.numpy.lcm must be integers.
exception error jax, numpy, dtype-validation, integer-required
The 'out' argument to jnp.compress is not supported.
exception error jax, unsupported-argument, numpy-compat, out-parameter
condition contains entries that are out of bounds
exception error jax, out-of-bounds, compress, length-mismatch
size must be positive and not greater than the size of the a
exception error jax, size-validation, out-of-range
y has more than 2 dimensions
exception error jax, statistics, covariance, shape-validation
m has more than 2 dimensions
exception error jax, statistics, covariance, ndim-validation
cov: dtype must be a subclass of float or complex; got {dtyp
exception error jax, covariance, dtype-validation
cannot handle multidimensional fweights
exception error jax, covariance, weights, shape-validation
incompatible numbers of samples and fweights
exception error jax, covariance, weights, length-mismatch
fweights must be integer.
exception error jax, covariance, weights, dtype-validation
cannot handle multidimensional aweights
exception error jax, covariance, weights, shape-validation
incompatible numbers of samples and aweights
exception error jax, covariance, weights, length-mismatch
corrcoef: dtype must be a subclass of float or complex; got
exception error jax, corrcoef, dtype-validation
with {nc} condition(s), either {nc} or {nc+1} functions are
exception error jax, piecewise, list-length-mismatch
{}-dimensional array given. Array must be at least two-dimen
exception error jax, linalg, matrix-power, ndim-validation
Last 2 dimensions of the array must be square
exception error jax, linalg, matrix-power, square-matrix-required
matrix_transpose requires at least 2 dimensions; got {ndim=}
exception error jax, numpy, shape-validation, linalg
Invalid order '{ord}' for vector norm.
exception error jax, linalg, norm, type-validation
ind must be a positive integer; got {ind=}
exception error jax, scipy, linalg, tensor, argument-validation
tensorinv is only possible when the product of the first `in
exception error jax, scipy, linalg, tensor, shape-validation
After moving axes to end, leading shape of a must match shap
exception error jax, numpy, linalg, tensor, shape-validation
Input arrays must have prod(a.shape[:b.ndim]) == prod(a.shap
exception error jax, numpy, linalg, tensor, shape-validation
multi_dot requires at least two arrays; got len(arrays)={len
exception error jax, numpy, linalg, matmul, argument-validation
multi_dot: input arrays must all be two-dimensional, except
exception error jax, numpy, linalg, matmul, shape-validation
multi_dot: last dimension of each array must match first dim
exception error jax, numpy, linalg, matmul, shape-mismatch
jnp.linalg.cond: input array must be at least 2D; got {arr.s
exception error jax, numpy, linalg, condition-number, shape-validation
jnp.linalg.cond: input array must not be empty; got {arr.sha
exception error jax, linalg, condition-number, empty-array, edge-case
jnp.linalg.cond: for {p=}, array must be square; got {arr.sh
exception error jax, numpy, linalg, condition-number, shape-validation
Input must be a rank-1 array.
validation error jax, numpy, polynomial, roots, shape-validation
expected deg >= 0
validation error jax, numpy, polynomial, polyfit, argument-validation
expected 1D vector for x
validation error jax, numpy, polynomial, polyfit, shape-validation
expected non-empty vector for x
validation error jax, numpy, polynomial, polyfit, empty-array, edge-case
expected 1D or 2D array for y
validation error jax, numpy, polynomial, polyfit, shape-validation
expected x and y to have same length
validation error jax, numpy, polynomial, polyfit, shape-mismatch
expected a 1-d array for weights
validation error jax, numpy, polynomial, polyfit, weights, shape-validation
expected w and y to have the same length
validation error jax, numpy, polynomial, polyfit, weights, length-mismatch
the number of data points must exceed order to scale the cov
validation error jax, polyfit, covariance, insufficient-data
input must be 1d or non-empty square 2d array.
validation error jax, numpy-poly, input-shape
Order of integral must be positive (see polyder)
validation error jax, polyint, argument-validation
k must be a scalar or a rank-1 array of length 1 or m.
validation error jax, polyint, shape-mismatch
Order of derivative must be positive
validation error jax, polyder, argument-validation
jnp.{name}: where must be None or a boolean array; got {wher
validation error jax, reductions, where-mask, dtype
The 'out' argument to jnp.{name} is not supported.
validation error jax, reductions, out-parameter, immutable-arrays
reduction operation {name} does not have an identity, so to
validation error jax, reductions, where-mask, identity
zero-size array to reduction operation {name} which has no i
validation error jax, reductions, empty-array
Named reductions not implemented for jnp.{name}()
validation error jax, named-axes, pmap, reductions
initial value must be a scalar. Got array of shape {initial_
validation error jax, reductions, initial-value, shape-mismatch
duplicate value in 'axis': {axis}
validation error jax, reductions, axis, duplicate
integer argument required; got dtype={arr.dtype}
validation error jax, reductions, dtype, internal
The 'out' argument to jnp.logaddexp.reduce is not supported.
validation error jax, logsumexp, out-parameter
The 'out' argument to jnp.logaddexp2.reduce is not supported
validation error jax, logaddexp2, out-parameter
The 'out' argument to jnp.mean is not supported.
validation error jax, mean, out-parameter
Axis must be specified when shapes of a and weights differ.
validation error jax, average, weights, shape-mismatch
Shape of weights must be consistent with shape of a along sp
validation error jax, average, weights, shape-mismatch
ddof and correction can't be provided simultaneously.
validation error jax, var, ddof, parameter-conflict
The 'out' argument to jnp.var is not supported.
validation error jax, var, out-parameter
jax.numpy.var does not yet support real dtype parameters whe
validation error jax, numpy, variance, complex-numbers, dtype
dtype argument to jnp.std must be inexact; got {dtype}
validation error jax, numpy, std, dtype-validation
The 'out' argument to jnp.std is not supported.
validation error jax, numpy, out-argument, immutable-arrays
The 'out' argument to jnp.ptp is not supported.
validation error jax, numpy, ptp, out-argument
The 'out' argument to jnp.nanmean is not supported.
validation error jax, numpy, nanmean, out-argument
The 'out' argument to jnp.nanvar is not supported.
validation error jax, numpy, nanvar, out-argument
The 'out' argument to jnp.nanstd is not supported.
validation error jax, numpy, nanstd, out-argument
The 'out' argument to jnp.{name} is not supported
validation error jax, numpy, cumsum, cumprod, out-argument
The input must be non-scalar to take a cumulative sum, howev
validation error jax, numpy, cumulative-sum, scalar-input
The input array has rank {x.ndim}, however axis was not set
validation error jax, numpy, cumulative-sum, axis-argument
The input must be non-scalar to take a cumulative product, h
validation error jax, numpy, cumulative-prod, scalar-input
jax.numpy.quantile does not support overwrite_input=True or
validation error jax, numpy, quantile, out-argument, overwrite-input
jax.numpy.nanquantile does not support overwrite_input=True
validation error jax, numpy, nanquantile, out-argument, overwrite-input
method can only be 'linear', 'lower', 'higher', 'midpoint',
validation error jax, numpy, quantile, invalid-argument
Weights cannot be complex types.
validation error jax, numpy, quantile, weights, complex-numbers
{method} doesn't support weights. Only method 'inverted_cdf'
validation error jax, numpy, quantile, weights, not-implemented
Weights shape must match 'a' shape when axis is None.
validation error jax, numpy, quantile, weights, shape-mismatch
Weights shape {weights.shape} must match reduction axes {tup
validation error jax, numpy, quantile, weights, shape-mismatch
'order' argument to argsort is not supported.
validation error jax, argsort, numpy-compat, unsupported-argument, sorting
jnp.partition for complex dtype is not implemented.
validation error jax, partition, complex-dtype, not-implemented, sorting
jnp.argpartition for complex dtype is not implemented.
validation error jax, argpartition, complex-dtype, not-implemented, sorting
need sequence of keys with len > 0 in lexsort
validation error jax, lexsort, empty-input, validation, sorting
all keys need to be the same shape
validation error jax, lexsort, shape-mismatch, validation, sorting
top_k is not compatible with complex inputs.
validation error jax, top-k, complex-dtype, value-error, sorting
mode must be 'largest' or 'smallest', got {mode!r}
validation error jax, top-k, invalid-argument, value-error, enum-validation
Number of tensordot axes (axes {}) exceeds input ranks ({} a
validation error jax, tensordot, shape-mismatch, numpy
tensordot requires axes lists to have equal length, got {} a
validation error jax, tensordot, axes-mismatch, numpy
tensordot requires both axes lists to be either ints, tuples
validation error jax, tensordot, type-error, axes
tensordot axes argument must be an int, a pair of ints, or a
validation error jax, tensordot, invalid-argument, numpy
The 'out' argument to jnp.outer is not supported.
validation error jax, outer, immutable-arrays, numpy-compat
out argument of {self}
validation error jax, ufunc, immutable-arrays, numpy-compat
where argument of {self}
validation error jax, ufunc, where-mask, numpy-compat
reduce only supported for binary ufuncs
validation error jax, ufunc, reduce, api-misuse
reduce only supported for functions returning a single value
validation error jax, ufunc, reduce, api-misuse
out argument of {self.__name__}.reduce()
validation error jax, ufunc, reduce, immutable-arrays
reduction operation {self.__name__!r} does not have an ident
validation error jax, ufunc, reduce, where-mask
where argument must have dtype=bool; got dtype={lax._dtype(w
validation error jax, ufunc, reduce, dtype, where-mask
tuple of axes
exception error jax, ufunc, reduce, multi-axis
zero-size array to reduction operation {self.__name__} which
exception error jax, ufunc, reduce, empty-array
accumulate only supported for binary ufuncs
exception error jax, ufunc, accumulate, api-misuse
accumulate only supported for functions returning a single v
exception error jax, ufunc, accumulate, api-misuse
out argument of {self.__name__}.accumulate()
exception error jax, ufunc, accumulate, immutable-arrays
accumulate does not allow multiple axes
exception error jax, ufunc, accumulate, axis
Because JAX arrays are immutable, jnp.ufunc.at() cannot oper
exception error jax, ufunc, scatter-add, immutable-arrays
reduceat only supported for binary ufuncs
exception error jax, ufunc, reduceat, api-misuse
num_segments must be non-negative.
validation error jax, segment-reduction, argument-validation, valueerror
unreduced for prod is not yet supported.
exception error jax, pallas-sharding, notimplementederror, distributed
unreduced for max is not yet supported.
exception error jax, sharding, notimplementederror, distributed
unreduced for min is not yet supported.
exception error jax, sharding, notimplementederror, distributed
Compiler params for platform {platform} cannot be used for {
validation error jax, pallas, backend-mismatch, gpu, tpu
MemoryRef type must be a ShapedArray, got {type(self.inner_a
validation error jax, pallas, aval, type-validation
Memory space {self.memory_space} is not supported by mesh {s
validation error jax, pallas, memory-space, mesh, gpu, tpu
Axis {axis} is out of bounds for grid {self.grid}
validation error jax, pallas, grid, index-out-of-bounds
Unsupported block dimension type: {type(dim)}. Allowed types
validation error jax, pallas, block-spec, type-validation
Unsupported block shape type: {type(dim)}
validation error jax, pallas, block-spec, internal
Array type must have a `shape` attribute, but got {type(arra
validation error jax, pallas, block-spec, aval
Block shape for {origin} (= {block_shape}) must have the sam
validation error jax, pallas, block-spec, shape-mismatch
shape polymorphism for Pallas does not support dynamically-s
validation error jax, pallas, shape-polymorphism, export, experimental
Index map function {debug_info.func_src_info} for {origin} m
validation error jax, pallas, block-spec, index-map
index_map returned a value of type {type(idx_aval)} at posit
validation error jax, pallas, index-map, block-spec
index_map returned a value of type {type(idx_aval)} at posit
validation error jax, pallas, index-map, scalar
Index map function {debug_info.func_src_info} for {origin} m
validation error jax, pallas, index-map, dtype
Index map function {debug_info.func_src_info} for {origin} m
validation error jax, pallas, index-map, closure
BlockMapping for {self.origin} has captured constants: {self
validation error jax, pallas, block-mapping, internal, closure
Cannot eagerly run with_memory_space_constraint.
validation error jax, pallas, eager-execution, primitive, memory-space
with_memory_space_constraint only supports arrays.
validation error jax, pallas, abstract-value, shapedarray, type-mismatch
When `check_vma=True` on `jax.shard_map`, `manual_axis_type`
validation error jax, shard-map, vma, shapedtypestruct, mesh, tpu, config-flag
Invalid out_shape type: {type(out_shape)}
validation error jax, pallas, out-shape, shapedtypestruct, input-validation
Unmatched parenthesis in {s!r}
validation error jax, einshape, einsum-notation, parse-error, parentheses
Equation must contain exactly one '->'
validation error jax, einshape, einsum-notation, equation-format, parse-error
Inconsistent size for {name}: {dim_sizes[name]} vs {shape_va
validation error jax, einshape, einsum-notation, dimension-mismatch, shape-validation
Size mismatch for group {group}: expected {shape_val}, got {
validation error jax, einshape, einsum-notation, shape-mismatch, grouped-reshape
get not supported yet
exception error jax, pallas, get, indexer, not-implemented
get not supported yet for block shape {b}
exception error jax, pallas, element-indexing, get, not-implemented
Concatenation with Element indexing is not yet supported.
exception error jax, pallas, concatenate, element-indexing, not-implemented
Shape along concat dimension {dimension} must be divisible b
validation error jax, pallas, concatenate, block-divisibility, valueerror
Stack with Element indexing is not yet supported.
exception error jax, pallas, stack, element-indexing, not-implemented
Stacking only supported when the block size along the stack
exception error jax, pallas, fusion, block-spec, shape-mismatch
Sum of sizes {n} must be equal to dimension {axis} of the op
exception error jax, pallas, lax-split, shape-mismatch, validation
Cannot permute last two dimensions with leading dimensions.
exception error jax, pallas, transpose, permutation, not-implemented
tile with non-int block dimensions not supported yet
exception error jax, pallas, tile, block-spec, not-implemented
Block size must be a multiple of the input size. Got block {
exception error jax, pallas, tile, shape-mismatch, divisibility
Every block dimension must be either a multiple or factor of
exception error jax, pallas, tile, divisibility, shape-mismatch
bitcast_convert_type with different bitwidths not supported
exception error jax, pallas, bitcast, dtype, not-implemented
Cannot pull iota along dimension {dimension} with None block
exception error jax, pallas, iota, block-spec, validation
reshape not supported yet: {aval_in}, {aval_out}
exception error
Expected {len(update_block_spec.block_shape)} start indices,
validation error
Parameter {key} is not a Jaxpr or sequence of Jaxprs: {value
validation error jax, pallas, jaxpr, control-flow, interpreter
run_scoped interpret rule does not support collective axes
exception error jax, pallas, collectives, interpreter, notimplemented
Padding with aliasing not supported.
exception error jax, pallas, aliasing, padding, interpreter
Acc ref must be at least 2D, got shape {shape}
validation error jax, pallas, tpu, mosaic, accumulator, shape
Acc ref dtype must be float32 or int32, got {dtype}
validation error jax, pallas, tpu, accumulator, dtype
Accumulators are not available on TPU {info.chip_version}
validation error jax, pallas, tpu, hardware, accumulator
mxu_id must be in [0, {info.num_mxus}), got {mxu_id=}
validation error jax, pallas, tpu, accumulator, index-out-of-range
The minor dimension size of an accumulator ref must be {info
validation error jax, pallas, tpu, accumulator, shape
The product of the major dimensions must be a multiple of {i
validation error jax, pallas, tpu, accumulator, shape, alignment
You can't use two different TensorCoreMeshes.
validation error jax, pallas, tpu, mesh, tensorcore, multi-core
cannot specify both devices and num_cores
validation error jax, pallas, tpu, mesh, mutually-exclusive-args
Unsupported core type: {core_type}
validation error jax, pallas, tpu, memory-space, lowering, version-skew
Invalid memory space: {memory_space!r}
validation error jax, pallas, tpu, memory-space, validation
More than one location found in string: {location_string}
validation error jax, pallas, tpu, error-handling, parsing
Could not find location in string {location_string}
validation warning jax, pallas, tpu, error-handling, parsing
Unknown action: {action}
validation error jax, pallas, tpu, async-copy, internal-api
Advanced indexers are not supported on TPU
validation error jax, pallas, tpu, interpreter, advanced-indexing
Revisited block {output_ranges[i]} of output {i} in iteratio
validation critical jax, pallas, tpu, interpreter, overlapping-writes, grid
Out-of-bounds read of ({device_id} {local_core_id} {memory_s
validation error jax, pallas, tpu, out-of-bounds, interpret-mode
Out-of-bounds block index {block_indices} for input "{input_
validation error jax, pallas, tpu, blockspec, out-of-bounds
Out-of-bounds write of ({device_id} {local_core_id} {memory_
validation error jax, pallas, tpu, out-of-bounds, write
Out-of-bounds block index {block_indices} for output "{outpu
validation error jax, pallas, tpu, blockspec, out-of-bounds
Out-of-bounds swap of ({device_id} {local_core_id} {memory_s
validation error jax, pallas, tpu, atomics, out-of-bounds
Out-of-bounds masked swap of ({device_id} {local_core_id} {m
validation error jax, pallas, tpu, atomics, mask, out-of-bounds
{primitive_name}: Buffers with a memory space of HBM or ANY
validation error jax, pallas, tpu, hbm, vmem, dma
masked load_p
validation error jax, pallas, tpu, interpret-mode, not-implemented, masked-load
run_scoped_p with collective axes is not supported
validation error jax, pallas, tpu, interpret-mode, not-implemented, collectives
Non-decrementing wait is not supported.
validation error jax, pallas, tpu, semaphore, interpret-mode, not-implemented
Unsupported block dim type: {type(b)}
validation error jax, pallas, blockspec, type-error
Dimensions with parallel semantics must form a prefix of the
validation error jax, pallas, tpu, grid, dimension-semantics
Unsupported memory space: {x.memory_space}
validation error jax, pallas, memory-space, not-implemented
Aliasing of scalar prefetch arguments is not currently suppo
validation error jax, pallas, tpu, aliasing, scalar-prefetch, interpret-mode
Kernel input {j} in HBM but does not have trivial BlockSpec.
validation error jax, pallas, tpu, hbm, blockspec, interpret-mode
Kernel output {j} in HBM but does not have trivial BlockSpec
validation error jax, pallas, tpu, hbm, blockspec, interpret-mode
Number of cores or threads must be at least 1, but got {self
validation error jax, pallas, tpu, config-validation, interpret-mode
Vector clock size must be at least 1, but got {self.vector_c
validation error jax, pallas, tpu, config-validation, race-detection
Vector clock size ({self.vector_clock_size}) must be greater
validation error jax, pallas, mosaic, tpu, interpret-mode, vector-clock, config
Logical shape {self.logical_shape} cannot be bigger than con
validation error jax, pallas, mosaic, shared-memory, shape-mismatch, interpret-mode
Range {rnge} is entirely out of bounds for shape {self.shape
validation error jax, pallas, mosaic, index-error, out-of-bounds, interpret-mode
Range {rnge} is (at least partially) out of bounds for alloc
exception error jax, pallas, mosaic, index-error, out-of-bounds, write
Attempting to deallocate allocation with key `{key}` that is
exception error jax, pallas, mosaic, shared-memory, key-collision, interpret-mode
Attempting to get contents of allocation with key `{key}` th
exception error jax, pallas, mosaic, shared-memory, key-collision, interpret-mode
Attempting to store into allocation with key `{key}` that is
exception error jax, pallas, mosaic, shared-memory, key-collision, interpret-mode
Attempting to swap into allocation with `key` {key} that is
exception error jax, pallas, mosaic, shared-memory, key-collision, swap
Semaphore {sem_id} occurs as both fixed-id and internal.
exception error jax, pallas, mosaic, semaphore, id-collision, interpret-mode
Computation failed on {self._failed_thread} with exception:\
exception error jax, pallas, mosaic, interpret-mode, threading, error-propagation
{uninitialized_memory} + {str(dtype)}
exception error jax, pallas, mosaic, dtype, not-implemented, interpret-mode
{axis} mixes JAX mesh and Pallas mesh grid axes
exception error jax, pallas, mesh, device-id, interpret-mode, not-implemented
{non_mesh_axes}
exception error jax, pallas, mesh, device-id, axis-name, interpret-mode
Unsupported device ID type: {device_id_type}
exception error jax, pallas, device-id, version-mismatch, interpret-mode
Unsupported range type: {type(r)}.
exception error jax, pallas, mosaic, indexing, unsupported-type
position {position} is out of range for clock {self.clock}
exception error jax, pallas, mosaic, vector-clock, race-detection, config
Too many dynamic shapes in the input. Mosaic currently only
validation error jax, pallas, mosaic, dynamic-shapes, limit-exceeded, lowering
program id was requested but no grid was provided.
exception error jax, pallas, mosaic, grid, program-id, lowering
Invalid axis {axis} for num_programs
exception error jax, pallas, mosaic, grid, axis, lowering
The Pallas TPU lowering currently supports in memory space A
validation error jax, pallas, tpu, block-mapping, memory-space
The Pallas TPU lowering currently requires that the last two
validation error jax, pallas, tpu, alignment, dynamic-shapes
The Pallas TPU lowering currently requires that rank 1 block
validation error jax, pallas, tpu, block-shape, rank-1
Pallas TPU requires a recent libtpu version (at least 0.0.46
exception critical jax, pallas, tpu, libtpu, version-mismatch
Non-trivial windowing is not supported for grid-free pallas_
exception error jax, pallas, tpu, grid, windowing
Unsupported block dimension type: {type(bd)} for block shape
exception error jax, pallas, tpu, block-shape, api-misuse
All block dimensions must be Elements or none of them can be
validation error jax, pallas, tpu, block-shape, element-window
Unsupported pipeline mode: {pipeline_mode}.
validation error jax, pallas, tpu, pipeline-mode, api-misuse
Lookahead is not supported for XLA pipeline emitter lowering
validation error jax, pallas, tpu, pipelining, lookahead
Only single (1) and double (2) buffering are supported. Got
validation error jax, pallas, tpu, pipelining, buffer-count
Cannot do int indexing on TPU
validation error jax, pallas, tpu, indexing
Loading from an accumulator is not supported. Use `matmul_po
validation error jax, pallas, tpu, accumulator, matmul
PRNG keys must be loaded from SMEM. Did you set the memory s
validation error jax, pallas, tpu, prng, random, smem
Indexing into a ()-shaped Ref not yet supported on TPU.
validation error jax, pallas, tpu, scalar-ref, smem
Can only load scalars from SMEM
validation error jax, pallas, tpu, smem, shape-mismatch
Loads are only allowed on VMEM and SMEM references.
validation error jax, pallas, tpu, memory-space, async-copy
Seed key_data must be 1D.
validation error jax, pallas, tpu, prng, shape-validation
Leading dimension of seed key_data must be 1.
validation error jax, pallas, tpu, prng, shape-validation
Can only load a single key per load.
validation error jax, pallas, tpu, prng, indexing
masked swap with non-32-bit data
validation error jax, pallas, tpu, masked-store, dtype
Expected value and mask to have the same shape, but got valu
validation error jax, pallas, tpu, shape-mismatch, masked-store
SMEM store does not support masks
validation error jax, pallas, tpu, smem, masked-store
Can only store scalars to SMEM
exception error jax, pallas, tpu, smem, shape-mismatch
Loads and stores are only allowed on VMEM and SMEM reference
exception error jax, pallas, tpu, memory-space, async-copy
Cannot store scalars to VMEM
exception error jax, pallas, tpu, vmem, scalar-store
Cannot swap scalars to VMEM.
exception error jax, pallas, tpu, swap, vmem, scalar
masked swap with strided store
exception error jax, pallas, tpu, strided-store, masked-store
Reductions over unsigned integers not implemented.
exception error jax, pallas, tpu, reduction, unsigned-integer, dtype
Reductions over {x_aval.dtype} not implemented.
exception error jax, pallas, tpu, reduction, dtype
`broadcast_to` is a Triton-specific primitive. Please consid
exception error jax, pallas, tpu, mosaic, broadcast
Unsigned integer dtype {aval.dtype} is not supported for dot
exception error jax, pallas, tpu, dot-general, unsigned-integer, dtype
{ctx.avals_out[0].dtype}
exception error jax, pallas, tpu, matmul, dtype, complex
Unsupported {preferred_element_type=}
exception error jax, pallas, tpu, dot-general, preferred-element-type, mixed-precision
Per-operand dot precision unsupported
exception error jax, pallas, tpu, dot-general, precision
Unsupported dot precision: {precision}
exception error jax, pallas, tpu, precision, dot-general
Per-operand conv precision unsupported
exception error jax, pallas, tpu, convolution, precision
Unsupported conv precision: {precision}
exception error jax, pallas, tpu, convolution, precision
Grouped convolutions are not supported on Pallas Mosaic TPU
exception error jax, pallas, tpu, convolution, grouped-conv, depthwise
Requires libtpu >= 0.1.0
exception error jax, pallas, tpu, libtpu, version, convolution
Unsigned integer dtype {aval.dtype} is not supported for con
exception error jax, pallas, tpu, convolution, unsigned-integer, dtype
Unsupported cast: {from_dtype} -> {to_dtype}
exception error jax, pallas, tpu, dtype-cast, unsigned-integer
Only arrays with 32-bit element types can be converted to sc
exception error jax, pallas, tpu, squeeze, scalar, dtype
Dimension must be 0 for 1D iota.
exception error jax, pallas, tpu, iota, indexing
Only take_along_axis-like gathers supported
exception error jax, pallas, tpu, gather, not-implemented
Only gathers along the two minormost dimensions supported on
exception error jax, pallas, tpu, tensorcore, gather
Unsupported gather
exception error jax, pallas, tpu, gather, not-implemented
{aval_out.dtype}
exception error jax, pallas, tpu, dtype, add
Unsupported constant type: {x.type}
exception error jax, pallas, mlir, constant-folding, internal
Unsupported combination of input dtype ({x_aval.dtype}) and
exception error jax, pallas, tpu, argmax, dtype
Only single axis reduction supported
exception error jax, pallas, tpu, argmax, reduction
Pallas top_k only supports float32 and bfloat16, got {input_
exception error jax, pallas, tpu, top-k, dtype
float32 top_k is not supported on TPUv3 or older
exception error jax, pallas, tpu, top-k, hardware-version
bfloat16 top_k is not supported on TPUv5 or older
exception error jax, pallas, tpu, top-k, hardware-version, bfloat16
Unsupported dtype for sign: {x.dtype}
exception error jax, pallas, tpu, dtype, sign
Not implemented: accuracy
exception error jax, pallas, tpu, rsqrt, accuracy
Transforms not supported for matmul_acc_lhs.
exception error jax, pallas, tpu, matmul, not-implemented
Transforms not supported for matmul_pop.
exception error jax, pallas, tpu, matmul, not-implemented
Element with {padding=} is not supported.
validation error jax, pallas, mosaic, pipeline, blockspec, padding
Unsupported block dimension type: {type(bd)}
exception error jax, pallas, blockspec, type-validation
next_fetch is None
exception error jax, pallas, internal-invariant, prefetch
Grid must be specified when using lookahead.
exception error jax, pallas, pipeline, lookahead, grid
Output buffering does not support lookahead.
exception error jax, pallas, pipeline, lookahead, output-buffer
Cannot partition grid over dynamic number of cores.
exception error jax, pallas, tpu, grid-partition, static-value
dimension_semantics must be the same length as grid.
exception error jax, pallas, dimension-semantics, validation
Cannot partition over cores without parallel grid dimensions
exception error jax, pallas, tpu, grid-partition, dimension-semantics
Cannot partition over multiple dynamic parallel dimensions:
exception error jax, pallas, grid-partition, dynamic-shape
Only one of src or dst can be a BufferedRef.
exception error jax, pallas, pipeline, copy
One of src or dst must be a BufferedRef.
exception error jax, pallas, pipeline, copy
Grid must consist of Python integers and JAX Arrays: {grid_t
exception error jax, pallas, grid, validation
All elements in the grid must be strictly positive, but got
exception error jax, pallas, grid, validation
Only one of `core_axis` or `core_axis_name` can be set.
exception error jax, pallas, mutually-exclusive-args, validation
The emit_pipeline body function must return None.
exception error jax, pallas, kernel-body, validation
WriteEffect should not apply to an input buffer {ref_invar_i
exception error jax, pallas, effects, input-mutation
WriteEffect on input buffer {ref_invar_idx}
exception error jax, pallas, effects, input-mutation
Grid mapping with hijax index maps are not currently support
exception error jax, pallas, hijax, index-map, not-implemented
Batching over custom allocations is not supported yet.
exception error jax, pallas, vmap, batching, allocations, not-implemented
Batching over dynamic grid values is not supported yet.
exception error jax, pallas, mosaic, vmap, batching, tpu, not-implemented
Not implemented: bitcast 1D
exception error jax, pallas, mosaic, bitcast, dtype, shape-validation
Not implemented: the 2nd minor dim can not be perfectly pack
exception error jax, pallas, mosaic, bitcast, dtype, packing, shape-validation
shift must be non-negative.
exception error jax, pallas, mosaic, roll, argument-validation, shift
axis is out of range.
exception error jax, pallas, mosaic, roll, axis-validation, shape-validation
stride and stride_axis must be both specified or not.
exception error jax, pallas, mosaic, roll, argument-validation, kwargs-pairing
stride must be non-negative.
exception error jax, pallas, mosaic, roll, argument-validation, stride
stride_axis is out of range
exception error jax, pallas, mosaic, roll, axis-validation, shape-validation
expected axis and stride_axis are different.
exception error jax, pallas, mosaic, roll, argument-validation
Either both or neither `src_sem` and `device_id` can be set.
exception error jax, pallas, mosaic, dma, remote-copy, argument-validation
Cannot `wait_send` on a local copy.
exception error jax, pallas, mosaic, dma, remote-copy, wait-send
dma_start not implemented in LoJAX yet.
exception error jax, pallas, mosaic, lojax, dma, not-implemented, tpu
DMA start priority must be non-negative: {priority}
exception error jax, pallas, mosaic, dma, priority, argument-validation
DMA source/destination/semaphore arguments must be Refs.
exception error jax, pallas, mosaic, dma, ref-vs-array, type-validation
Cannot signal on a non-()-shaped semaphore: {dst_sem_shape}
exception error jax, pallas, mosaic, dma, semaphore, shape-validation
DMA source semaphore must be a Ref.
exception error jax, pallas, mosaic, dma, semaphore, ref-vs-array, type-validation
Cannot signal on a non-()-shaped semaphore: {src_sem_shape}
exception error jax, pallas, mosaic, dma, semaphore, shape-validation
DMA partial discharge add=True not yet implemented.
exception error jax, pallas, mosaic, dma, discharge, not-implemented, pipelining
The number of sources must match the packing factor ({packin
exception error jax, pallas, tpu, dtype, shape-mismatch
The bitwidth of `x` must match the bitwidth of `unpacked_dty
exception error jax, pallas, tpu, dtype, bitwidth
Index {index} is out of bounds for packing factor {packing_f
exception error jax, pallas, tpu, index-out-of-bounds
with_memory_space_constraint only supports HBM, VMEM, SMEM,
exception error jax, pallas, tpu, memory-space, unsupported-operation
trace_value requires a scalar value, got shape {value.shape}
exception error jax, pallas, tpu, debugging, scalar-required
trace_value requires i32 or f32, got {value.dtype}
exception error jax, pallas, tpu, debugging, dtype
load_staged_rhs must be an integer or None.
exception error jax, pallas, tpu, matmul, type-error
Expected an accumulator ref, got {acc}
exception error jax, pallas, tpu, matmul, memory-space
The shape of the accumulator {acc_shape} is not 2-dimensiona
exception error jax, pallas, tpu, matmul, shape-mismatch
The shape of the accumulator {acc_shape} does not match the
exception error jax, pallas, tpu, matmul, shape-mismatch
Only float32 and int32 results are supported, got {dtype}
exception error jax, pallas, tpu, matmul, dtype
Bit width must be 32
exception error jax, pallas, tpu, random, prng
Cannot split a Pallas key. Use fold_in instead to generate n
exception error jax, pallas, tpu, random, prng, not-implemented
block_size ({len(block_size)}) and tile_size ({len(tile_size
exception error jax, pallas, tpu, random, shape-mismatch
The current TPU does not have SparseCores
exception critical jax, pallas, tpu, sparsecore, hardware, environment
Mesh has {self.num_cores} cores, but the current TPU chip ha
exception error jax, pallas, tpu, sparsecore, hardware, resource-limits
You can't use two different ScalarSubcoreMeshes.
exception error jax, pallas, tpu, sparsecore, mesh-compat
{self} should have the same core axis name and number of cor
validation error jax, pallas, tpu, sparsecore, mesh-compat
{self} should have a different axis name from the TensorCore
validation error jax, pallas, tpu, sparsecore, axis-naming
Mesh has {self.num_subcores} subcores, but the current TPU c
validation error jax, pallas, tpu, hardware-mismatch, sparsecore
You can't use two different VectorSubcoreMeshes.
validation error jax, pallas, tpu, sparsecore, mesh-compat
{self} should have the same core axis name and number of cor
validation error jax, pallas, tpu, sparsecore, mesh-compat
{self} should have a different core axis name from the Tenso
validation error jax, pallas, tpu, sparsecore, axis-naming
{self} should have a different subcore axis name from the Te
validation error jax, pallas, tpu, sparsecore, axis-naming
Unsupported dtype: {dtype}
validation error jax, pallas, tpu, sparsecore, dtype
Get does not support loading from {ref_memory_space!r}. Copy
exception error jax, pallas, tpu, sparsecore, memory-space
Integer indexing of refs that follows a non-trivial slice is
exception error jax, pallas, tpu, sparsecore, indexing
Get only supports slices with stride 1, got {strides}
exception error jax, pallas, tpu, sparsecore, slicing, strides
Get only supports loading scalars from SMEM.
exception error jax, pallas, tpu, sparsecore, scalar-load
Get does not support masked scalar loads
exception error jax, pallas, tpu, sparsecore, masked-load
Swap does not support storing to {ref_memory_space!r}. Copy
exception error jax, pallas, tpu, sparsecore, memory-space, swap
Swap only supports slices with stride 1, got {strides}
exception error jax, pallas, tpu, sparsecore, swap, memory-space
Swap only supports scalars in SMEM.
exception error jax, pallas, tpu, sparsecore, masked-store, swap
Swap does not support masked scalar stores
exception error jax, pallas, tpu, sparsecore, atomic, swap
Gather only supports loading from VMEM, got {ref_aval.memory
exception error jax, pallas, tpu, sparsecore, memory-space, gather
{x.shape=} does not match expected shape {expected_shape}
exception error jax, pallas, sparsecore, scatter, shape-mismatch
val.dtype={x.dtype} != ref.dtype={ref.dtype}
exception error jax, pallas, sparsecore, scatter, dtype-mismatch
{mask.shape=} does not match expected shape {expected_shape}
exception error jax, pallas, sparsecore, scatter, mask-shape
Scatter only supports VectorSubcoreMesh, got {type(ref_aval.
exception error jax, pallas, sparsecore, mesh, scatter
Scatter only supports storing to VMEM, got {memory_space}
exception error jax, pallas, sparsecore, scatter, memory-space
Indices must not be empty
exception error jax, pallas, sparsecore, scatter, empty-indices
Cannot bitcast a ()-shaped array to a dtype with a different
exception error jax, pallas, sparsecore, bitcast, scalar
Cannot bitcast from {x.dtype} ({old_bitwidth} bits) to {dtyp
exception error jax, pallas, sparsecore, bitcast, divisibility
x.dtype={x.dtype} must be uint32, int32 or float32
exception error jax, pallas, sparsecore, scan-count, dtype
mask.dtype={mask.dtype} is not a boolean dtype
exception error jax, pallas, sparsecore, scan-count, mask-dtype
x.shape={x.shape} != mask.shape={mask.shape}
exception error jax, pallas, sparsecore, scan-count, shape-mismatch
reductions require axes to be (0,) on SparseCore, but got {a
exception error jax, pallas, sparsecore, reduction, axes
cummax: x={x.aval} must be rank 1
exception error jax, pallas, sparsecore, cummax, rank
packed cannot be specified if layout is specified.
validation error jax, pallas, mosaic-gpu, tmem, layout, api-misuse
packed, collective and layout arguments are only supported f
validation error jax, pallas, mosaic-gpu, smem, gmem, tmem, api-misuse
Cannot specify both out_shape and out_type. Use out_type.
exception error jax, pallas, mosaic-gpu, deprecation, api-migration
Cannot specify both scratch_shapes and scratch_types. Use sc
exception error jax, pallas, mosaic-gpu, deprecation, scratch-memory
Unsupported dtype: {ref.dtype}
exception error jax, pallas, mosaic-gpu, dtype, shared-memory, not-implemented
Only byte-aligned shapes are supported. Got shape: {ref.dtyp
exception error jax, pallas, mosaic-gpu, shared-memory, alignment, shape-validation
Only SMEM and TMEM refs are supported.
exception error jax, pallas, mosaic-gpu, ref-union, memory-space, not-implemented
Ref unions can't be assigned to.
exception error jax, pallas, mosaic-gpu, ref-union, immutable, read-only
layout attribute is only defined for TMEM refs
exception error jax, pallas, mosaic-gpu, tmem, gpu-kernels
collective attribute is only defined for TMEM refs
exception error jax, pallas, mosaic-gpu, tmem, collective-ops
Some aliased TMEM references are collective and some are not
exception error jax, pallas, mosaic-gpu, tmem, aliasing
All aliased Refs must have the same memory space (SMEM or TM
exception error jax, pallas, mosaic-gpu, smem, tmem, aliasing
Cannot transform type: {x}
exception error jax, pallas, abstract-values, transforms, type-dispatch
Transpose cannot be moved before a tiling transform when it
exception error jax, pallas, mosaic-gpu, transpose, tiling, transforms
Strided slices unsupported. Got stride: {ds.stride}
exception error jax, pallas, slicing, strides, tiling
Expected slice start ({start}) and slice size ({size}) to be
exception error jax, pallas, alignment, tiling, slicing
Unsupported index type: {type(idx)}
exception error jax, pallas, indexing, type-dispatch, tiling
Commuting a `UntilingTransform` with a `ReshapeTransform` is
exception error jax, pallas, mosaic-gpu, reshape, transforms
Commuting a `UntilingTransform` with a `ReshapeTransform` is
exception error jax, pallas, mosaic-gpu, tiling, transforms
Unfolding dimensions is not supported when commuting an `Un
exception error jax, pallas, mosaic-gpu, reshape, unfolding
Folding tiled dimensions into untiled dimensions is not supp
exception error jax, pallas, mosaic-gpu, reshape, tiling
Cannot commute `UntilingTransform` with `ReshapeTransform` w
exception error jax, pallas, mosaic-gpu, reshape, tiling, layout
Unsupported transform: {type(transform)}
exception error jax, pallas, mosaic-gpu, transforms, not-implemented
ref must be a reference
exception error jax, pallas, mosaic-gpu, peer-memory, type-error
Can't make a multicast reference into a peer reference.
exception error jax, pallas, mosaic-gpu, peer-memory, multicast
dims and idxs must have the same length
exception error jax, pallas, mosaic-gpu, cluster, validation
Can't make a peer reference into a multicast reference.
exception error jax, pallas, mosaic-gpu, multicast, peer-memory
Can't transpose a TMEM reference.
exception error jax, pallas, mosaic-gpu, tmem, transpose
Unsupported type: {x}
exception error jax, pallas, mosaic-gpu, aval, type-error
Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
exception error jax, pallas, mosaic-gpu, swizzle, validation
Swizzle {self.swizzle} requires the trailing dimension to be
exception error jax, pallas, mosaic-gpu, swizzle, layout, shape-mismatch
Can't transpose the swizzled dimension.
exception error jax, pallas, mosaic-gpu, transpose, swizzle
Reshape shape {shape} is not divisible by swizzle elements {
validation error jax, pallas, reshape, swizzle, mosaic-gpu
Cannot commute unswizzle and indexer with {aval}, which does
validation error jax, pallas, dtype, swizzle, indexing
Non-slice indices are not supported in 2 minormost dims: {id
validation error jax, pallas, indexing, swizzle, not-implemented
Swizzled dims cannot be sliced
validation error jax, pallas, slicing, swizzle
Unsupported ndim: {x.ndim}
validation error jax, pallas, ndim, shape-validation
Unsupported shape: {x.shape}
validation error jax, pallas, shape-validation, ndim
Second dimension {x.shape[1]} must be divisible by batch_siz
validation error jax, pallas, shape-validation, divisibility
Expected indexer to have exactly {k + 2} dimensions, but got
validation error jax, pallas, indexing, rank-mismatch
Slicing batch dimensions is not supported.
validation error jax, pallas, indexing, slicing, not-implemented
collective_axes is not supported in pallas_call. Use plgpu.k
validation error jax, pallas, mosaic-gpu, collective-axes, unsupported-feature
Num arrivals must be at least 1, but got {n}
validation error jax, pallas, barrier, validation, config
Barriers are not arrays
validation error jax, pallas, barrier, api-misuse
Cluster barriers are not arrays
validation error jax, pallas, barrier, cluster, api-misuse
Preinitialized WGMMAAccumulatorRef only supported in pl.run_
validation error jax, pallas, wgmma, accumulator, api-misuse
Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
validation error jax, pallas, mosaic-gpu, wgmma, not-implemented, indexing
cluster= must be at most 3D, got {self}.
validation error jax, pallas, mosaic-gpu, cluster, launch-config, validation
grid_names must have the same length as grid, got {self}.
validation error jax, pallas, mosaic-gpu, launch-config, validation, shape-mismatch
cluster_names must have the same length as cluster, got {sel
validation error jax, pallas, mosaic-gpu, launch-config, validation, shape-mismatch
num_threads and thread_name must be either both set or both
validation error jax, pallas, mosaic-gpu, warp-specialization, validation
Requested too many CUDA threads per block. Each Mosaic threa
validation error jax, pallas, mosaic-gpu, thread-limits, resource-limit, validation
Can't instantiate {self} with arguments.
validation error jax, pallas, mosaic-gpu, layout, api-misuse, validation
Only TiledLayout supports reductions.
validation error jax, pallas, mosaic-gpu, layout, reduction, unsupported-type
Invalid TMEMLayout: {self}
validation error jax, pallas, mosaic-gpu, tmem-layout, enum, version-mismatch
{tiling=} and {grid=} must have same length.
validation error jax, pallas, mosaic-gpu, nd-loop, tiling, shape-mismatch
Tiling {tiling} does not divide grid {grid}.
validation error jax, pallas, mosaic-gpu, nd-loop, tiling, divisibility
Expected metadata dtype to be uint2, got: {meta.dtype}
validation error jax, pallas, mosaic-gpu, sparse, tcgen05, dtype, validation
Expected metadata to be 3-dimensional (M, K // 4, 2), but it
validation error jax, pallas, mosaic-gpu, sparse, tcgen05, shape, rank
Expected the trailing dimension of the metadata to be 2, got
validation error jax, pallas, mosaic-gpu, sparse, tcgen05, shape, last-dim
Sparse metadata format not implemented for {operand_dtype=}
validation error jax, pallas, mosaic-gpu, sparse, tcgen05, dtype, not-implemented
No valid out swizzle{what}: minor dimension has {minor_dim_b
validation error jax, pallas, mosaic-gpu, swizzle, shared-memory, alignment, matmul
`thread_id` must be zero when allocating a buffer for all th
exception error jax, pallas, mosaic-gpu, interpret-mode, allocation, warp-specialization
`block_id` must be zero when allocating a buffer for all thr
exception error jax, pallas, mosaic-gpu, interpret-mode, allocation, block-id
Out-of-bounds read of {allocation_key}: reading [{read_range
exception error pallas, mosaic-gpu, interpret-mode, out-of-bounds, jax
Out-of-bounds block index {block_indices} for {allocation_ke
exception error pallas, mosaic-gpu, interpret-mode, out-of-bounds, input-block
Dynamic indexing not supported in GPU interpret mode
exception error pallas, mosaic-gpu, interpret-mode, dynamic-indexing
Unsupported transform: {transform}
exception error pallas, mosaic-gpu, interpret-mode, unsupported-transform, internal
Out-of-bounds swap of {allocation_key}: swapping [{read_writ
exception error pallas, mosaic-gpu, interpret-mode, out-of-bounds, atomics, swap
Out-of-bounds masked swap of {allocation_key}: swapping [{re
exception error pallas, mosaic-gpu, interpret-mode, out-of-bounds, atomics, mask
Unsupported barrier type: {type(barrier)}
exception error pallas, mosaic-gpu, interpret-mode, barrier, unsupported
predicate not supported
exception error pallas, mosaic-gpu, interpret-mode, predicate, not-implemented
reduction_op not supported
exception error pallas, mosaic-gpu, interpret-mode, reduction, not-implemented
tcgen05_mma only allows arriving on a Barrier
exception error pallas, mosaic-gpu, tcgen05, mma, barrier, interpret-mode
tcgen05_mma only allows arriving on a Barrier that orders te
exception error pallas, mosaic-gpu, tcgen05, mma, barrier, tensor-core
tcgen05_commit_arrive only allows arriving on a Barrier
exception error pallas, mosaic-gpu, tcgen05, barrier, interpret-mode
tcgen05_commit_arrive only allows arriving on a Barrier that
exception error pallas, mosaic-gpu, tcgen05, barrier, tensor-core
Dynamic grid bounds not (yet) supported in GPU interpret mod
exception error pallas, mosaic-gpu, interpret-mode, dynamic-grid, not-implemented
Unsupported mesh type: {type(mesh)}
exception error pallas, mosaic-gpu, mesh, interpret-mode, type-validation
Invalid grid {grid_mapping.grid} in grid_mapping: expected g
exception error pallas, mosaic-gpu, grid, mesh, config-mismatch
Dynamic grid bounds not (yet) supported on GPU
exception error pallas, mosaic-gpu, interpret-mode, dynamic-grid, not-implemented
Scalar arguments not (yet) supported on GPU
exception error pallas, mosaic-gpu, interpret-mode, scalar-args, not-implemented
Unsupported memory space: {space}
exception error pallas, mosaic-gpu, memory-space, interpret-mode, not-implemented
Requesting collective allocations, but no explicit thread ax
exception error pallas, mosaic-gpu, collective, mesh, thread-axis, config
Collective allocations along cluster axes are not supported.
exception error jax, mosaic-gpu, interpret-mode, collective-allocations
Scoped allocation must have the thread axis in its collectiv
exception error jax, mosaic-gpu, run-scoped, collective-axes
Indexing barrier with {transforms} not supported in GPU inte
exception error jax, mosaic-gpu, barrier, transforms
Expected an `NDIndexer`, but got {transforms[0]}
exception error jax, mosaic-gpu, barrier, type-validation
Attempting to operate on barrier without indexing, but `num_
exception error jax, mosaic-gpu, barrier, indexing
Dynamic grid bounds not supported in the Mosaic GPU lowering
exception error
Data type bitcast is only supported from i8 to other types.
exception error jax, pallas, mosaic-gpu, bitcast, dtype, not-implemented
Only workgroup memory is supported but got {ref}.
exception error jax, pallas, mosaic-gpu, smem, memory-space, bitcast
Data type bitcast is only supported for contiguous 1D arrays
exception error jax, pallas, mosaic-gpu, bitcast, stride, contiguity
Can not bitcast memory region of size {shape_bitwidth} bits
exception error jax, pallas, mosaic-gpu, bitcast, size-mismatch, alignment
TMEM aliasing only supported for Refs with the same first di
exception error jax, pallas, tmem, tcgen05, aliasing, shape-mismatch
Only byte-aligned bitcasts are supported.
exception error jax, pallas, mosaic-gpu, smem, bitcast, alignment, sub-byte
The base ref for aliases must come from a slice_smem op.
exception error jax, pallas, mosaic-gpu, warpgroup, smem, aliasing, not-implemented
Unsupported TMEM ref {ref}.
exception error jax, pallas, tmem, tcgen05, aliasing, not-implemented
Unsupported memory space.
exception error jax, pallas, mosaic-gpu, memory-space, aliasing, not-implemented
Tiling without swizzle is not supported.
exception error
Transforms are not yet implemented for warpgroup semantics
exception error jax, pallas, mosaic-gpu, transforms, warpgroup
Stores to TMEM are asynchronous operations and cannot be per
exception error jax, pallas, mosaic-gpu, tmem, blackwell, tcgen05
Can only store scalars in warp-level lowering.
exception error jax, pallas, mosaic-gpu, warp-semantics, store
Can only store to references (got {x_ref}).
exception error jax, pallas, mosaic-gpu, type-error, internal-api
Can only store scalars or vectors (got {value}).
exception error jax, pallas, mosaic-gpu, type-error, warpgroup
Can only store to references (got {x_smem}).
exception error jax, pallas, mosaic-gpu, type-error, warpgroup
out_dtype argument in binary_op_lowering_rule_wg
exception error jax, pallas, mosaic-gpu, dtype, lowering, not-implemented
{ctx.prim} does not support {x_aval.dtype} and {y_aval.dtype
exception error jax, pallas, mosaic-gpu, dtype, binary-op, not-implemented
Unsupported dtype {x_aval.dtype}
exception error jax, pallas, mosaic-gpu, integer-pow, complex-dtype, not-implemented
Not implemented: accuracy
exception error jax, pallas, mosaic-gpu, rsqrt, accuracy, not-implemented
No axis names are available. Make sure you are using `pl.cor
exception error jax, pallas, mosaic-gpu, cluster, mesh, core-map
Unknown cluster axis {axis_name}, available axes: {[*axis_na
exception error jax, pallas, mosaic-gpu, cluster, axis-name, mesh
No axis names are available. Make sure you are using `pl.cor
exception error jax, pallas, mesh, collective, axis-name
Axis {axis_name} does not refer to a GPU mesh axis (availabl
exception error jax, pallas, mesh, collective, axis-name
Ordered debug_print is not supported on Pallas.
exception error jax, pallas, debug-print, ordered
Only positional arguments are supported by debug_print on Pa
exception error jax, pallas, debug-print, kwargs
Can only print scalars in warp-level code.
exception error jax, pallas, debug-print, warp-level
debug_print only supports printing of scalar values, or a si
exception error jax, pallas, debug-print, arrays
Buffer poisoning is not supported on GPU yet.
exception error jax, pallas, buffer-poisoning, debug-flag, gpu
Only thread-collective allocations are supported in run_scop
exception error jax, pallas, run-scoped, collective-axes, warp-specialization
WGMMA accumulators can only be allocated non-collectively. H
exception error jax, pallas, wgmma, run-scoped, collective-axes
Invalid WGMMA accumulator dtype for s8/i8 WGMMA. Expected si
exception error jax, pallas, wgmma, int8, dtype
Only thread-collective allocations are supported in multithr
exception error jax, pallas, run-scoped, warp-specialization, collective-axes
Can't convert to ref: {aval}
exception error jax, pallas, run-scoped, ref, type-error
get_global only supports semaphores, got {what}
exception error jax, pallas, get-global, semaphore
Uninitialized Refs are not supported in lowering of run_stat
exception error jax, pallas, run-state, uninitialized-ref
Integer constant out of range for i64: {v}
exception error pallas, mosaic-gpu, integer-overflow, jax
Unsupported index: {v} of type {type(v)}
exception error pallas, mosaic-gpu, indexing, type-error
Cannot merge empty list of indexers
exception error pallas, internal-api, indexing
Non-unit strides not implemented.
exception error pallas, mosaic-gpu, slicing, not-implemented
Unhandled transforms for semaphore_read: {transforms}
exception error pallas, mosaic-gpu, semaphore, not-implemented
Non-debug checks are not supported by the Mosaic GPU backend
exception error pallas, checkify, mosaic-gpu, not-implemented
relayout_p is not supported with Lane semantics.
exception error pallas, mosaic-gpu, layout, not-implemented
Only plgpu.Layout.WG_SPLAT is supported for scalar values.
exception error pallas, mosaic-gpu, layout, scalar
Failed to infer the output layout of the iota. Please apply
exception error pallas, mosaic-gpu, iota, layout
Encountered non-leading UntilingTransform or UnswizzleRef tr
exception error pallas, mosaic-gpu, transforms, not-implemented
dynamic grid bounds not supported in the Mosaic GPU backend
exception error pallas, mosaic-gpu, grid, not-implemented
Mesh {mesh} is not supported by the Mosaic GPU backend
exception error pallas, mosaic-gpu, mesh, not-implemented
Lowering multiple mesh/function pairs is not supported by th
exception error pallas, mosaic-gpu, mpmd-map, not-implemented
External meshes are not supported by the Mosaic GPU backend
exception error pallas, mosaic-gpu, mpmd-map, mesh, not-implemented
Scratch operands are not supported by the Mosaic GPU mpmd_ma
exception error pallas, mosaic-gpu, mpmd-map, scratch, not-implemented
Unsupported block size type: {type(bd)}
exception error pallas, mosaic-gpu, blockspec, not-implemented
Only collective_axes that include all JAX device mesh axes
exception error jax, pallas, mosaic-gpu, collective-axes, device-mesh, multicast
Non-indexing transforms on GMEM refs are not implemented.
exception error jax, pallas, mosaic-gpu, gmem-ref, transforms, not-implemented
The stored value has shape {src.shape}, but the target refer
exception error jax, pallas, mosaic-gpu, shape-mismatch, async-store
The stored value has dtype {src.dtype}, but the target refer
exception error jax, pallas, mosaic-gpu, dtype-mismatch, async-store
Expected a single barrier, got a barrier reference with shap
exception error jax, pallas, mosaic-gpu, barrier, async-store
Scalars are not supported in async_store_smem
exception error jax, pallas, mosaic-gpu, scalar, async-store, not-implemented
Unexpected unhandled transforms: {remaining_ref_transforms}
exception error jax, pallas, mosaic-gpu, internal-invariant, transforms, async-store
async_store_smem requires a tiled and swizzled ref
exception error jax, pallas, mosaic-gpu, tiled-layout, swizzle, async-store
Can only transfer integer bytes (shape={shape}, dtype={dtype
exception error jax, pallas, mosaic-gpu, sub-byte-dtype, alignment, async-store
Transfer is not a multiple of {WARPGROUP_SIZE} bytes
exception error jax, pallas, mosaic-gpu, alignment, warpgroup, async-store
Expected source shape to be {expected_src_shape}, but got {s
exception error jax, pallas, mosaic-gpu, shape-mismatch, sharding, gmem-copy
Can only transfer integer bytes (shape={dst_ty.shape}, dtype
exception error jax, pallas, mosaic-gpu, sub-byte-dtype, alignment, gmem-copy
Expected exactly one collective axis, got {collective_axes=}
exception error jax, pallas, mosaic-gpu, collective-axes, leader-tracked
Partitioned loads only supported for clusters of size 2. Got
exception error jax, pallas, mosaic-gpu, cluster-size, leader-tracked, not-implemented
copy_gmem_to_smem with a barrier is only supported Hopper an
exception error jax, pallas, mosaic-gpu, gpu-architecture, tma, barrier, cp-async
Only the TMA implementation supports collective copies
exception error jax, pallas, mosaic-gpu, gpu-architecture, collective-axes, cp-async
Only the TMA implementation supports leader_tracked copies
exception error jax, pallas, mosaic-gpu, gpu-architecture, leader-tracked, cp-async
The cp.async implementation only supports oob_mode=OOBFillMo
exception error jax, pallas, mosaic-gpu, gpu-architecture, out-of-bounds, cp-async
The cp.async implementation does not support user-defined pr
exception error jax, pallas, mosaic-gpu, gpu-architecture, predicate, cp-async, not-implemented
copy_gmem_to_smem without a barrier is only supported on pre
exception error mosaic-gpu, pallas, gpu, barrier, tma, hopper
Only copies transferring a number of bytes divisible by the
exception error mosaic-gpu, pallas, alignment, warpgroup, copy
predicate is not supported with Warpgroup lowering in jaxlib
exception error mosaic-gpu, pallas, jaxlib-version, predicate, tma
Loading from a remote ref is only supported in jaxlib versio
exception error mosaic-gpu, pallas, jaxlib-version, multi-gpu, peer-id
`collective_axes` must be specified when `leader_tracked` is
validation error mosaic-gpu, pallas, api-misuse, collectives
wait_gmem_to_smem is only supported on pre-Hopper GPUs, whic
validation error mosaic-gpu, pallas, gpu, hopper, synchronization
GMEM refs with peer ids are not supported in warpgroup lower
validation error mosaic-gpu, pallas, prefetch, multi-gpu, peer-id
Barrier does not support slice with `stride != 1`
validation error mosaic-gpu, pallas, barrier, indexing, slice
Barrier can only be indexed with integers or slices, got {id
validation error mosaic-gpu, pallas, barrier, indexing, type-error
Barrier does not support arbitrary transforms
validation error mosaic-gpu, pallas, barrier, transforms
Arriving on a collective barrier is not supported in a warp
validation error mosaic-gpu, pallas, barrier, warp-specialization
Incompatible shapes for matrix multiplication: lhs={a.shape}
validation error mosaic-gpu, pallas, wgmma, shape-mismatch, matrix-multiply
Expected WGMMAAbstractAccumulatorRef got {acc_aval}
validation error mosaic-gpu, pallas, wgmma, accumulator, type-error
Accumulator shape {inner.shape} does not match value shape {
validation error jax, pallas, mosaic-gpu, wgmma, shape-mismatch
Accumulator dtype {inner.dtype} does not match value dtype {
validation error jax, pallas, mosaic-gpu, wgmma, dtype-mismatch
Accumulator and LHS have incompatible shapes. Expected LHS t
validation error jax, pallas, tcgen05, mma, shape-mismatch
Accumulator and RHS have incompatible shapes. Expected RHS t
validation error jax, pallas, tcgen05, collective-mma, shape-mismatch
Accumulator and RHS have incompatible shapes. Expected RHS t
validation error jax, pallas, tcgen05, mma, shape-mismatch
a_scale and b_scale must both be present or absent.
validation error jax, pallas, tcgen05, block-scaling, argument-validation
Accumulator must be a TMEM Ref.
validation error jax, pallas, tcgen05, tmem, memory-space
LHS must be a TMEM/SMEM Ref.
validation error jax, pallas, tcgen05, smem, tmem, memory-space
RHS must be an SMEM Ref.
validation error jax, pallas, tcgen05, smem, memory-space
Accumulator Ref must be collective if collective_axis is set
validation error jax, pallas, tcgen05, collective-mma, tmem
LHS Ref must be collective if collective_axis is set.
validation error jax, pallas, tcgen05, collective-mma, tmem
MMA barrier must have orders_tensor_core set to True.
validation error jax, pallas, tcgen05, barrier, pipelining
a_scale must be a TMEM Ref
validation error jax, pallas, tcgen05, block-scaling, tmem
b_scale must be a TMEM Ref
validation error jax, pallas, tcgen05, block-scaling, tmem
Unsupported transforms for ACC: {acc_transforms}.
exception error jax, pallas, tcgen05, transform, not-implemented
Unsupported transforms for LHS: {a_transforms}.
exception error jax, pallas, tcgen05, transform, not-implemented
MMA lhs tiling does not fit swizzle. {lhs_tiling=} expected=
validation error jax, pallas, tcgen05, swizzle, tiling, layout
Unsupported transforms for RHS: {b_transforms}.
exception error jax, pallas, tcgen05, transform, not-implemented
MMA rhs tiling does not fit swizzle {rhs_tiling=} expected={
validation error jax, pallas, tcgen05, swizzle, tiling, layout
MMA rhs swizzle must match lhs swizzle. {lhs_swizzle=} {rhs_
validation error jax, pallas, mosaic-gpu, tcgen05, mma, swizzle
TMEM transpose not allowed.
validation error jax, pallas, tcgen05, tmem, transpose
Unsupported transforms: {a_scale_transforms}
exception error jax, pallas, tcgen05, fp8, scaling, not-implemented
Unsupported transforms: {b_scale_transforms}
exception error jax, pallas, tcgen05, fp8, scaling, not-implemented
Unsupported transforms: {a_sparse_metadata_transforms}
exception error jax, pallas, tcgen05, sparse, not-implemented
Unsupported transforms for {ref}. Transforms {transforms}.
exception error jax, pallas, tcgen05, warp-group, transforms, not-implemented
No cluster axes found.
validation error jax, pallas, tcgen05, cluster, collective-mma
Can only perform collective MMA along minormost cluster axis
validation error jax, pallas, tcgen05, cluster, collective-mma, axis-order
inline_mgpu_p only supports plgpu.ShapeDtypeStruct return ty
validation error jax, pallas, mosaic-gpu, inline-mgpu, type-validation
inline_mgpu_p only supports only SomeLayout and RefType arg
validation error jax, pallas, mosaic-gpu, inline-mgpu, type-validation
Mismatched type shape: {treedef} != {treedef_ty}
validation error jax, pallas, inline-mgpu, pytree, signature-mismatch
Mismatched type: {a, t}
validation error jax, pallas, inline-mgpu, type-mismatch
inline_mgpu_p does not support discharge.
exception error jax, pallas, inline-mgpu, discharge, not-implemented
Array dtype mismatch: expected {v.mlir_dtype} got {mlir_dtyp
validation error jax, pallas, inline-mgpu, dtype-mismatch
Array shape mismatch: expected {ty.shape} got {v.shape}.
validation error jax, pallas, inline-mgpu, shape-mismatch
Array layout mismatch: expected {v.layout} got {ty.layout.to
validation error jax, pallas, inline-mgpu, layout-mismatch
Unexpected layout for {v} (expected: {ty})
validation error jax, pallas, inline-mgpu, layout-mismatch
Unexpected type {ty} for value {v}
validation error jax, pallas, inline-mgpu, type-validation
inline_mgpu in a single-warp context only supports scalar ar
validation error jax, pallas, inline-mgpu, warp-semantics, scalar-required
Not all transforms could be handled. Remaining transforms: {
exception error jax, pallas, mosaic-gpu, transforms, not-implemented
Transform mismatch: got {user_transforms}, expected {ty_tran
validation error jax, pallas, transform-mismatch, validation
inline_mgpu in a single-warp context only supports scalar re
validation error jax, pallas, warp-semantics, scalar-return
inline_mgpu_p return type tree mismatch: {ret} != {return_ty
validation error jax, pallas, pytree, return-type-mismatch
Unsupported aval type: {aval}, {type(aval)}, {t}
validation error jax, pallas, aval, not-implemented, mlir
Expected a ShapeDtypeStruct, but got: {r}
validation error jax, pallas, shapedtypestruct, return-type
Expected a FragmentedArray, but got: {fa}
validation error jax, pallas, fragmented-array, lowering
Expected {result_shape} but got {fa.shape}
validation error jax, pallas, shape-mismatch, lowering
Output layout {out_layout} does not match the layout of the
validation error jax, pallas, layout-mismatch, mosaic-gpu
Expected 0D shape, but got {fa.shape}
validation error jax, pallas, scalar, shape-mismatch
Expected WGSplatFragLayout, but got {fa.layout}
validation error jax, pallas, layout, scalar, mosaic-gpu
Async load only supports TMEM refs
validation error jax, pallas, tmem, memory-space, blackwell
Unsupported dtype for reduction: {val_aval.dtype}. Only floa
validation error jax, pallas, tmem, dtype, reduction
Unimplemented transforms for TMEM refs. {transforms=}
validation error jax, pallas, tmem, transforms, not-implemented
Unsupported load reduce operation: {orig_reduce}
validation error jax, pallas, tmem, reduction, invalid-argument
Async store only supports TMEM refs
validation error jax, pallas, tmem, memory-space, blackwell
Only collective_axes that include all JAX device mesh ({mesh
exception error jax, pallas, tpu, distributed, device-mesh, multimem
Unhandled transforms for multimem_store: {transforms}
exception error jax, pallas, transforms, multimem, lowering
JAX device mesh is required by multimem_load_reduce, but not
exception error jax, pallas, device-mesh, multimem, missing-context
Failed to infer the output layout of multimem_load_reduce. P
exception error jax, pallas, layout, multimem, tpu
Only tiled and WG strided layouts are supported by multimem_
validation error jax, pallas, layout, multimem, unsupported-layout
Unhandled transforms for multimem_load_reduce: {transforms}
exception error jax, pallas, transforms, multimem, lowering
Cannot pass the same ref into a mpmd map multiple times
exception error jax, pallas, mpmd, ref-aliasing, notimplementederror
MPMD map with more than one mesh requires scratch_type to ha
exception error jax, pallas, mpmd, tpu, scratch-memory, multi-mesh
The kernel function in mpmd_map {debug_info.func_src_info} s
exception error jax, pallas, mpmd, kernel-api, inplace-semantics
input pinned buffers without input_output_aliases:{missing}
exception error jax, pallas, pallas-call, input-output-alias, pinned-buffer
pallas_call requires all mesh axes to be Manual, got {get_ab
exception error jax, pallas, pallas-call, sharding, mesh, manual-axes
pallas_call does not support hijax for index_map
exception error jax, pallas, pallas-call, index-map, lowering, notimplementederror
interpret with dynamic grid bounds unsupported
exception error jax, pallas, jvp, autodiff, dynamic-grid, notimplementederror
JVP with aliasing not supported.
exception error jax, pallas, jvp, autodiff, input-output-alias, notimplementederror
pallas_call with a mesh does not support JVP
exception error jax, pallas, jvp, autodiff, mesh, spmd, notimplementederror
vmapping pallas_call with no arguments.
exception error jax, pallas, vmap, batching, not-implemented
pallas_call with a mesh does not support batching
exception error jax, pallas, vmap, mesh, sharding, not-implemented
The kernel function in the pallas_call {debug_info.func_src_
exception error jax, pallas, closure, constants, tracing
The kernel function in the pallas_call {debug_info.func_src_
exception error jax, pallas, kernel, return-value, validation
Cannot lower a pallas_call with constants.
exception error jax, pallas, lowering, constants, compiler
Only interpret mode is supported on CPU backend.
exception error jax, pallas, cpu, backend, interpret-mode
Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
exception error jax, pallas, rocm, amd, gpu, backend, mosaic
The format string expects {n_placeholders} argument{'' if n_
exception error pallas, format-string, debug-print, jax
run_scoped discharge does not support collective_axes yet.
exception error pallas, run-scoped, collectives, not-implemented, jax
Cannot handle new consts created by state discharge.
exception error pallas, state-discharge, run-scoped, not-implemented, jax
run_scoped lowering outside of Pallas does not support colle
exception error pallas, run-scoped, collectives, jax
get_global discharge is not supported in interpret mode.
exception error pallas, get-global, interpret-mode, not-implemented, jax
Cannot {name} on a non-semaphore Ref: {sem_aval}
exception error pallas, semaphore, type-validation, jax
Cannot {name} on a non-()-shaped semaphore: {sem_shape}
exception error pallas, semaphore, shape-validation, jax
Must {name} semaphores of the following types: {allowed_sema
exception error pallas, semaphore, dtype-validation, jax
Semaphore value shape {ref_value.shape} does not match aval
exception error pallas, semaphore, shape-mismatch, state-discharge, jax
Must signal an int32 value, but got {value_aval.dtype}
exception error pallas, semaphore, dtype-validation, jax
`device_id`s must be an int32 value, but got {aval.dtype}
exception error pallas, semaphore, dtype-validation, device-id, jax
Remote signal not implemented.
exception error pallas, semaphore, remote-signal, not-implemented, jax
Multiple core support not implemented.
exception error pallas, semaphore, multi-core, not-implemented, jax
Must wait an int32 value.
exception error pallas, semaphore, dtype-validation, jax
{axis_name} mixes JAX mesh and Pallas mesh grid axes
exception error pallas, mesh, collectives, axis-name, not-implemented, jax
`device_id_type` must be MESH if `device_id` is a dict, got:
exception error pallas, device-id, validation, jax
Number of device ids must match the number of mesh axes, but
exception error pallas, device-id, mesh, validation, jax
Unsupported device id type: {device_id_type}
exception error pallas, device-id, not-implemented, jax
TransformedRef must have been indexed before passing into ja
exception error pallas, refs, jaxpr-call, validation, jax
Only positional arguments are supported by debug_print on Pa
validation error jax, pallas, triton, debug-print, kwargs, gpu
unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}
validation error jax, pallas, triton, jnp-minimum, dtype, notimplementederror
end must be greater than start, but got: {end} <= {start}
validation error jax, pallas, triton, arange, iota, empty-range, validation
start and end must fit in int32
validation error jax, pallas, triton, int32-overflow, arange, large-tensors
cannot splat a tensor
validation error jax, pallas, triton, splat, scalar-shape, reshape, typeerror
cannot cast {src} tp {dst_type}
validation error jax, pallas, triton, dtype-cast, float8, notimplementederror
cannot cast from `{dtype_name}`
validation error jax, triton, pallas, dtype, cast, gpu
cannot cast to `{dtype_name}`
validation error jax, triton, pallas, dtype, cast, gpu
cannot cast {src} to {dst_type}
exception error jax, triton, pallas, mlir, cast, not-implemented
Only 2-argument concatenate is supported.
exception error jax, triton, pallas, concatenate, not-implemented
Only concatenate along the last dimension is supported.
exception error jax, triton, pallas, concatenate, axis, not-implemented
Only arguments with shape [..., 1] are supported.
exception error jax, triton, pallas, concatenate, shape, not-implemented
Only 2-argument stack is supported in Triton.
exception error jax, triton, pallas, stack, not-implemented
Only stack along the last dimension is supported in Triton.
exception error jax, triton, pallas, stack, axis, not-implemented
Only unstack of size 2 is supported in Triton.
exception error jax, triton, pallas, unstack, not-implemented
Only unstack along the last dimension is supported in Triton
exception error jax, triton, pallas, unstack, axis, not-implemented
Only power-of-2 num parts supported.
exception error jax, triton, pallas, split, not-implemented
Only equal-sized splits are supported.
exception error jax, triton, pallas, split, shape, not-implemented
Unexpected dim_block_size: {dim_block_size}
exception error jax, triton, pallas, internal, block-mapping, version-mismatch
unsupported cache modifier: {cache_modifier}
exception error jax, triton, pallas, load, cache-modifier, validation
unsupported eviction policy: {eviction_policy}
exception error jax, triton, pallas, load, eviction-policy, validation
loading from a block pointer is not supported
exception error jax, triton, pallas, load, block-pointer, not-implemented
unsupported pointer type: {ptr_type}
exception error jax, triton, pallas, load, pointer, validation
other requires mask to be provided
exception error jax, triton, pallas, load, mask, validation
other cannot be a block if pointer is not a block
exception error jax, triton, pallas, load, shape, validation
mask cannot be a block if pointer is not a block
exception error jax, triton, pallas, load, mask, shape, validation
Unsigned integer dtype {aval.dtype} is not supported for dot
exception error jax, pallas, triton, gpu, matmul, dtype
Unsupported dot algorithm: {precision}.
exception error jax, pallas, triton, precision, dot-algorithm
Unsupported dot precision: {precision}.
exception error jax, pallas, triton, precision, validation
a and b must be 2D, but got: {a_type.shape} and {b_type.shap
exception error jax, pallas, triton, shape, matmul
float64 dot requires M>=16, N>=8, K>=16 per warp tile (Trito
exception error jax, pallas, triton, float64, tile-shape
a and b must have the same element type, but got: {a_type.el
exception error jax, pallas, triton, dtype, matmul
Reductions with constants not supported.
exception error jax, pallas, triton, reduction, closures
`index_type` must be i32.
exception error jax, pallas, triton, argmax, index-dtype
`pallas` reduce operations only support one reduce axis.
exception error jax, pallas, triton, reduction, axes
Axis name {axis_name} not found in grid.
exception error jax, pallas, triton, grid, axis-index
Only single indexer is supported.
exception error jax, pallas, triton, atomics, indexing
unsupported atomic operation: {atomic_type}
exception error jax, pallas, triton, atomics
cmp and val must have identical dtypes and shapes
validation error jax, pallas, triton, dtype, atomic-cas
ref must be scalar.
validation error jax, pallas, triton, atomic-cas, shape
cmp must be scalar.
validation error jax, pallas, triton, atomic-cas, shape
val must be scalar.
validation error jax, pallas, triton, atomic-cas, shape
`next_power_of_2` requires a non-negative integer.
validation error jax, pallas, utils, validation
Scan with {num_extensive_outputs} extensive output(s) is not
validation error jax, pallas, triton, scan, control-flow
Scan with {num_extensive_inputs} extensive argument(s) is no
validation error jax, pallas, triton, scan, control-flow
erf_inv_lowering_helper not implemented for {x.dtype}
exception error jax, pallas, triton, dtype, erfinv
sign_lowering_helper not implemented for {x.dtype}
exception error jax, pallas, triton, dtype, sign
The two inputs to `nextafter` must have the same dtype, but
validation error jax, pallas, triton, dtype, nextafter
`nextafter` only supports float32 and float64, but got {x.dt
validation error jax, pallas, triton, dtype, nextafter
unreduced cannot contain None. All elements in unreduced sho
validation error jax, sharding, partition-spec, mesh
reduced cannot contain None. All elements in reduced should
validation error jax, sharding, partition-spec, mesh
`unreduced` and `reduced` argument to PartitionSpec cannot o
validation error jax, sharding, partition-spec, mesh
Expected unreduced_kind to be of type `jax.sharding.Unreduce
validation error jax, sharding, partition-spec, type-error
`unreduced_kind` should be `None` when `unreduced` is an emp
validation error jax, sharding, partition-spec, validation
partitions cannot overlap with unreduced axes passed to Part
validation error jax, sharding, partition-spec, mesh
partitions cannot overlap with reduced axes passed to Partit
validation error jax, sharding, partition-spec, mesh
Missing required keyword argument: 'in_sharding'
validation error jax, sharding, missing-argument, pjit
Context mesh cannot be empty. Please use `jax.set_mesh` API
validation error jax, mesh, sharding, context
layouts passed to `with_layout_constraint` must be of type `
validation error jax, layout, type-validation
with_layout_constraint in eager mode can only be applied to
validation error jax, layout, eager-mode
All args passed to `explicit_layout` must have the same type
validation error jax, layout, consistency
Missing required keyword argument: 'in_layouts'
validation error jax, layout, missing-argument
'devices' argument to pmap must be non-empty, or None.
validation error jax, pmap, devices, empty-argument
pmap in_axes must be an int, None, or (nested) container wit
validation error jax, pmap, in-axes, type-validation
pmap out_axes must be an int, None, or (nested) container wi
validation error jax, pmap, out-axes, type-validation
pmap requires at least one argument with a mapped axis.
validation error jax, pmap, in-axes, axis-size
pmapped function has static_broadcasted_argnums={static_broa
validation error jax, pmap, static-argnums, argument-count
pmap in_axes: value not a valid prefix (dynamic message incl
validation error jax, pmap, in-axes, pytree-prefix
Specified axis_size {axis_size} doesn't match received axis_
validation error jax, pmap, axis-size, shape-mismatch
Only one profiler server can be active at a time.
validation error jax, profiler, singleton
No active profiler server.
validation error jax, profiler, lifecycle
Profile has already been started. Only one profile may be ru
error_code error jax, profiler, concurrency, singleton
Invalid trace folder: {latest_trace_folder}
validation warning jax, profiler, perfetto, filesystem
No profile started
error_code error jax, profiler, lifecycle
dirichlet requires alpha.ndim >= 1, got alpha.ndim == {}
validation error jax, random, dirichlet, input-validation, shape
dtype argument to `exponential` must be a float dtype, got {
validation error jax, random, exponential, dtype, input-validation
method argument to `gamma` must be one of {'exact', 'approxi
validation error jax, random, gamma, method, input-validation
dtype argument to `gamma` must be a float dtype, got {dtype}
validation error jax, random, gamma, dtype, input-validation
method argument to `loggamma` must be one of {'exact', 'appr
validation error jax, random, loggamma, method, input-validation
method argument to `poisson` must be one of {'exact', 'appro
validation error jax, random, poisson, method, input-validation
lam shape must be broadcastable to shape argument; got lam.s
validation error jax, random, poisson, broadcasting, shape
`poisson` with method='exact' is only implemented for the th
validation error jax, random, poisson, prng, not-implemented, tpu
dtype argument to `gumbel` must be a float dtype, got {dtype
validation error jax, random, gumbel, dtype, input-validation
Must provide valid mode for gumbel got: %s
validation error jax, random, gumbel, mode, input-validation
Number of samples without replacement ({k}) cannot exceed nu
validation error jax, random, categorical, sampling, top-k, shape
dtype argument to `laplace` must be a float dtype, got {dtyp
validation error jax, random, laplace, dtype, input-validation
dtype argument to `logistic` must be a float dtype, got {dty
validation error jax, random, logistic, dtype, input-validation
dtype argument to `pareto` must be a float dtype, got {dtype
validation error jax, random, pareto, dtype, input-validation
dtype argument to `t` must be a float dtype, got {dtype}
validation error jax, random, t-distribution, dtype, input-validation
method argument to `chisquare` must be one of {'exact', 'app
validation error jax, random, chisquare, method, input-validation
dtype argument to `chisquare` must be a float dtype, got {dt
validation error jax, random, chisquare, dtype, input-validation
dtype argument to `f` must be a float dtype, got {dtype}
validation error jax, random, f-distribution, dtype, input-validation
dtype argument to `maxwell` must be a float dtype, got {dtyp
validation error jax, random, maxwell, dtype, input-validation
dtype argument to `double_sided_maxwell` must be a float dty
validation error jax, random, dtype-validation
dtype argument to `weibull_min` must be a float dtype, got {
validation error jax, random, dtype-validation
dtype argument to `rayleigh` must be a float dtype, got {dty
validation error jax, random, dtype-validation
dtype argument to `wald` must be a float dtype, got {dtype}
validation error jax, random, dtype-validation
dtype argument to `geometric` must be an int dtype, got {dty
validation error jax, random, dtype-validation, discrete
dtype argument to `triangular` must be a float dtype, got {d
validation error jax, random, dtype-validation
dtype argument to `lognormal` must be a float or complex dty
validation error jax, random, dtype-validation
dtype argument to `binomial` must be a float dtype, got {dty
validation error jax, random, dtype-validation
{name} requires all arguments to have matching type. Got key
validation error jax, sharding, shard-map, pvary, random
PRNG key seed must be a scalar; got {seed!r}.
exception error jax, prng, seed, philox
PRNG key seed must be an integer; got {seed!r}
exception error jax, prng, seed, dtype-validation
philox2x32_random_bits got invalid prng key.
exception error jax, prng, philox, key-validation
requires 8-, 16-, 32- or 64-bit field width.
exception error jax, prng, philox, argument-validation
random bits array of size exceeding 2 ** 64
exception error jax, prng, philox, shape-validation, resource-limits
PRNG key seed must be a scalar; got {seed!r}.
exception error jax, prng, seed, philox
PRNG key seed must be an integer; got {seed!r}
exception error jax, prng, seed, dtype-validation
philox4x32_random_bits got invalid prng key.
exception error jax, prng, philox, key-validation
requires 8-, 16-, 32- or 64-bit field width.
exception error jax, prng, philox, argument-validation
random bits array of size exceeding 2 ** 64
exception error jax, prng, philox, shape-validation, resource-limits
PRNG with name {impl.name} already registered: {impl}
exception error jax, prng, registry, duplicate-registration
JAX encountered invalid PRNG key data: expected key_data to
exception error jax, prng, typeerror, input-validation
JAX encountered invalid PRNG key data: expected key_data.ndi
exception error jax, prng, ndim, typeerror
JAX encountered invalid PRNG key data: expected key_data.sha
exception error jax, prng, shape-mismatch, typeerror
JAX encountered invalid PRNG key data: expected key_data.dty
exception error jax, prng, dtype, uint32
iteration over a 0-d key array
exception error jax, prng, iteration, unpacking
key array cannot be converted to boolean.
exception error jax, prng, truthiness, bool
JAX array with PRNGKey dtype cannot be converted to a NumPy
exception error jax, prng, numpy, serialization
random_unwrap takes key array operand, got {keys.dtype=}
exception error jax, prng, unwrap, dtype
_rbg_random_bits got invalid prng key.
exception error jax, prng, rbg, shape-mismatch
requires 8-, 16-, 32- or 64-bit field width.
exception error jax, prng, bit-width, dtype
Expected base_key to be a typed PRNG key; got {self._base_ke
exception error jax, prng, stateful-rng, typed-key
Expected counter to be a scalar integer ref; got {self._coun
exception error jax, prng, stateful-rng, ref, experimental
cannot operate on split stateful generator
exception error jax, prng, stateful-rng, split, spawn
When used within transformed code, jax.experimental.random.s
exception error jax, prng, stateful-rng, jit, tracing
PRNG key seed must be a scalar; got {seed!r}.
exception error jax, prng, seed, scalar, shape
PRNG key seed must be an integer; got {seed!r}
exception error jax, prng, seed, dtype, integer
threefry_2x32 requires uint32 arguments, got {}
exception error jax, prng, threefry, uint32, dtype
threefry_random_bits got invalid prng key.
exception error jax, prng, threefry, key-validation
requires 8-, 16-, 32- or 64-bit field width.
exception error jax, prng, threefry, bit-width, dtype
Only the type=1 case of eigh is implemented.
validation error jax, linalg, eigenvalue, not-implemented
Only the eigvals=None case of eigh is implemented.
validation error jax, linalg, eigenvalue, not-implemented
Expected 'output' to be either 'real' or 'complex', got {out
validation error jax, linalg, schur, argument-validation
Unsupported QR decomposition mode '{mode}'
validation error jax, linalg, qr-decomposition, argument-validation
mode must be 'right' or 'left', got {mode!r}
validation error jax, linalg, qr-decomposition, argument-validation
Array shapes are not compatible for Q @ c operation: a has s
validation error jax, linalg, shape-mismatch, qr-decomposition
Array shapes are not compatible for c @ Q operation: a has s
validation error jax, linalg, shape-mismatch, qr-decomposition
Expected assume_a to be one of {valid_assume_a}; got {assume
validation error jax, linalg, solve, argument-validation, version-migration
Invalid 'trans' value {trans}
validation error jax, linalg, solve-triangular, argument-validation
Expected A to be a (batched) square matrix, got {A.shape=}.
validation error jax, linalg, matrix-exponential, shape-mismatch
expected A to be a square matrix
validation error jax, linalg, internal-api, shape-mismatch
A.dtype={A.dtype} is not supported.
validation error jax, linalg, dtype, matrix-exponential
expected A to be a (batched) square matrix, got A.shape={A_a
validation error jax, linalg, frechet-derivative, shape-mismatch
expected E to be a (batched) square matrix, got E.shape={E_a
validation error jax, linalg, frechet-derivative, shape-mismatch
expected A and E to be the same shape, got A.shape={A_arr.sh
validation error jax, linalg, frechet-derivative, shape-mismatch
Arguments to jax.scipy.linalg.block_diag must have at most 2
validation error jax, linalg, block-diagonal, shape-mismatch
diagonal and off-diagonal values must have same dtype, got {
validation error jax, linalg, tridiagonal, dtype
Only float32 and float64 inputs are supported as inputs to j
validation error jax, linalg, tridiagonal, dtype
for select='i', select_range must be specified.
validation error jax, linalg, tridiagonal, missing-argument
Got empty index range in select_range.
validation error jax, linalg, tridiagonal, argument-validation
eigh_tridiagonal(..., select='v') is not implemented
validation error jax, eigenvalues, not-implemented, linalg
'select must have a value in {'a', 'i', 'v'}.
exception error jax, argument-validation, linalg, eigenvalues
method='qdwh' only supports mxn matrices where m < n where s
exception error jax, polar-decomposition, qdwh, shape-mismatch
Unknown polar decomposition method {method}.
exception error jax, argument-validation, polar-decomposition
The input `a` must be at least a 2-D array.
exception error jax, polar-decomposition, input-validation, shape-error
The argument `side` must be either 'right' or 'left'.
exception error jax, argument-validation, polar-decomposition
Blocked version is not implemented yet.
exception warning jax, sqrtm, not-implemented, linalg
Input 'T' must be square.
exception error jax, rsf2csf, schur-form, shape-validation
Input 'Z' must be square.
exception error jax, rsf2csf, schur-form, shape-validation
Input array shapes must match: Z: {Z_arr.shape} vs. T: {T_ar
exception error jax, rsf2csf, shape-mismatch
hankel: c must be at least 1-dimensional, got a scalar.
exception error jax, hankel, input-validation, shape-error
hankel: r must be at least 1-dimensional, got a scalar.
exception error jax, hankel, input-validation, shape-error
The length of f along the last axis must be at least 2; got
exception error jax, leslie-matrix, input-validation
Incorrect lengths for f and s. The length of s along the las
exception error jax, leslie-matrix, off-by-one, argument-length-validation
The length of `a` along the last axis must be at least 2; go
exception error jax, companion-matrix, polynomial, input-validation
fiedler_companion requires the last axis of 'a' to have nonz
exception error jax, fiedler-companion, polynomial, empty-array
n must be a positive integer; got {n}.
exception error jax, convolution-matrix, argument-validation
convolution_matrix: a must be at least 1-dimensional, got a
exception error jax, convolution-matrix, input-validation, shape-error
len(a) must be at least 1; got shape {a_arr.shape}.
exception error jax, convolution-matrix, empty-array
mode must be one of 'full', 'valid', 'same'; got {mode!r}.
exception error jax, convolution-matrix, argument-validation, invalid-enum-argument
Expected kind to be on of: {valid_kind}; got {kind}
exception error jax, scipy, linalg, argument-validation
Expected kind to be one of: {valid_kind}; got {kind}
exception error jax, scipy, linalg, argument-validation
n must be a positive power of 2; got {n}.
exception error jax, scipy, linalg, matrix, argument-validation
scale must be None, 'sqrtn', or 'n'; got {scale!r}.
exception error jax, scipy, linalg, dft, argument-validation
dtype must be a complex floating-point type; got {dtype}.
exception error jax, scipy, linalg, dtype, dft
Unrecognized method {method}. The two valid methods are eith
exception error jax, scipy, linalg, sylvester, argument-validation
Incompatible shapes for Sylvester equation: A: {A.shape} B:
exception error jax, scipy, linalg, shape-mismatch, sylvester
coordinates must be a sequence of length input.ndim, but {}
exception error jax, scipy, ndimage, coordinates, shape-mismatch
jax.scipy.ndimage.map_coordinates does not yet support mode
exception error jax, scipy, ndimage, unsupported-feature, boundary-mode
jax.scipy.ndimage.map_coordinates currently requires order<=
exception error jax, scipy, ndimage, interpolation, unsupported-feature
args argument to jax.scipy.optimize.minimize must be a tuple
exception error jax, scipy, optimize, type-error, argument-validation
Method {method} not recognized
exception error jax, scipy, optimize, argument-validation
in1 and in2 should have the same dimensionality
exception error jax, scipy, signal, fft, shape-mismatch
mode must be one of ['same', 'full', 'valid']
exception error jax, scipy, signal, argument-validation
mapped axes must have same shape; got {in1.shape=} {in2.shap
exception error jax, scipy, signal, vmap, shape-mismatch
For 'valid' mode, One input must be at least as large as the
exception error jax, scipy, signal, convolution, shape-mismatch
Unrecognized {mode=}
exception error jax, scipy, signal, internal-api, argument-validation
mode must be one of ['full', 'same', 'valid']
exception error jax, scipy, signal, convolution, argument-validation
in1 and in2 must have the same number of dimensions
exception error jax, scipy, signal, convolution, shape-mismatch
zero-size arrays not supported in convolutions, got shapes {
exception error jax, scipy, convolution, empty-array, shape-validation
One input must be smaller than the other in every dimension.
exception error jax, scipy, convolution, shape-mismatch
unsupported mode: {mode}
exception error jax, scipy, convolution, invalid-argument-value
Got {method=}; expected 'auto', 'fft', or 'direct'.
exception error jax, scipy, convolution, invalid-argument-value
convolve2d() only supports boundary='fill', fillvalue=0
exception error jax, scipy, convolution, not-implemented, boundary
convolve2d() only supports 2-dimensional inputs.
exception error jax, scipy, convolution, rank-mismatch, shape-validation
correlate2d() only supports boundary='fill', fillvalue=0
exception error jax, scipy, correlation, not-implemented, boundary
correlate2d() only supports 2-dimensional inputs.
exception error jax, scipy, correlation, rank-mismatch, shape-validation
overwrite_data argument not implemented.
exception error jax, scipy, detrend, not-implemented, immutable-arrays
Trend type must be 'linear' or 'constant'.
exception error jax, scipy, detrend, invalid-argument-value
Breakpoints must be non-negative and less than length of dat
exception error jax, scipy, detrend, breakpoints, index-out-of-range
The extension length n ({n}) is too big. It must not exceed
exception error jax, scipy, stft, signal-extension, shape-validation
Unknown value for mode {mode}, must be one of: ('psd', 'stft
exception error jax, scipy, internal-api, spectral-analysis
Unknown boundary option '{boundary}', must be one of: {list(
exception error jax, scipy, stft, boundary, invalid-argument-value
two-argument mode is available only when mode=='psd'
exception error jax, scipy, internal-api, spectral-analysis
two-arguments must have the same rank ({x.ndim} vs {y.ndim})
exception error jax, scipy, csd, rank-mismatch, shape-validation
x and y cannot be broadcast together.
exception error jax, scipy, csd, broadcasting, shape-mismatch
nperseg must be a positive integer
exception error jax, scipy, stft, nperseg, jit-tracer, argument-validation
nfft must be greater than or equal to nperseg.
exception error jax, scipy, stft, fft-size, argument-validation
noverlap must be less than nperseg.
exception error jax, scipy, stft, overlap, argument-validation
betainc does not support complex-valued inputs.
validation error jax, scipy-special, complex-dtype, dtype-validation
dawsn does not support complex-valued inputs.
validation error jax, scipy-special, dawson, complex-dtype
entr does not support complex-valued inputs.
validation error jax, scipy-special, entropy, complex-dtype
kl_div does not support complex-valued inputs.
validation error jax, scipy-special, kl-divergence, complex-dtype
rel_entr does not support complex-valued inputs.
validation error jax, scipy-special, relative-entropy, complex-dtype
Riemann zeta function not implemented; pass q != None to com
exception error jax, scipy-special, zeta, not-implemented
Argument `n` to polygamma must be of integer type. Got dtype
validation error jax, scipy-special, polygamma, dtype-validation
polygamma does not support complex-valued inputs.
validation error jax, scipy-special, polygamma, complex-dtype
x.dtype={} is not supported, see docstring for supported typ
exception error jax, scipy-special, normal-cdf, dtype-validation
invalid value ({e.ty}) encountered in ndtri.
exception error jax, scipy-special, nan, ndtri, floating-point
series_order must be a Python integer.
exception error jax, scipy-special, log-ndtr, type-validation
series_order must be non-negative.
validation error jax, scipy-special, log-ndtr, argument-validation
series_order must be <= 30.
validation error jax, scipy-special, log-ndtr, argument-validation
x.dtype={np.dtype(dtype)} is not supported.
exception error jax, scipy-special, log-ndtr, dtype-validation
complex input not supported.
validation error jax, scipy-special, bessel, complex-dtype
Negative orders for normalization is not implemented yet.
exception error jax, scipy-special, legendre, not-implemented, spherical-harmonics
z.dtype={} is not supported, see docstring for supported typ
exception error jax, scipy-special, legendre, dtype-validation
z must be a 1D array.
validation error jax, scipy-special, legendre, not-implemented, argument-validation
Computations for m!=n are not yet supported.
exception error jax, scipy, legendre, not-implemented, argument-mismatch
The 'diff_n' argument to jax.scipy.special.sph_harm_y is not
exception error jax, scipy, spherical-harmonics, not-implemented, version-mismatch
expi does not support complex-valued inputs.
exception error jax, expi, complex-numbers, unsupported-dtype
Argument `x` to sici must be real-valued. Got dtype {x.dtype
exception error jax, sici, complex-numbers, unsupported-dtype
expn does not support complex-valued inputs.
exception error jax, expn, complex-numbers, unsupported-dtype, exponential-integral
exp1 does not support complex-valued inputs.
exception error jax, exp1, complex-numbers, unsupported-dtype, exponential-integral
x.dtype={dtype} is not supported, see docstring for supporte
exception error jax, spence, dtype, typeerror, dilogarithm
n must be a non-negative integer.
exception error jax, scipy, bernoulli, static-argument, input-validation
jnp.poch does not support complex-valued inputs.
exception error jax, scipy, poch, complex-dtype, unsupported-operation
hyp1f1 does not support complex-valued inputs.
exception error jax, scipy, hyp1f1, complex-dtype, unsupported-operation
Illegal nan_policy value {nan_policy!r}; expected one of {'p
exception error jax, scipy, stats, mode, nan-policy, argument-validation
Logic for `nan_policy` of {nan_policy} is not implemented
exception error jax, scipy, stats, mode, not-implemented, nan-policy
In order to best JIT compile `mode`, we cannot know whether
exception error jax, scipy, stats, mode, jit, not-implemented, nan-policy
In order to best JIT compile `rankdata`, we cannot know whet
exception error jax, scipy, stats, rankdata, jit, not-implemented, nan-policy
unknown method '{method}'
exception error jax, scipy, stats, rankdata, argument-validation, method
{nan_policy} is not supported
exception error jax, scipy, sem, nan-handling
`alpha` must be one-dimensional; got alpha.shape={alpha.shap
exception error jax, scipy, dirichlet, shape-validation
`x` must have either the same number of entries as `alpha` o
exception error jax, scipy, dirichlet, shape-validation
gaussian_kde does not support complex data
exception error jax, scipy, kde, complex-dtype
`dataset` input should have multiple elements.
exception error jax, scipy, kde, input-validation
`weights` input should be one-dimensional.
exception error jax, scipy, kde, weights, shape-validation
`weights` input should be of length n
exception error jax, scipy, kde, weights, length-mismatch
`bw_method` should be 'scott', 'silverman', a scalar, or a c
exception error jax, scipy, kde, bandwidth, argument-validation
mean does not have dimension {self.d}
exception error jax, scipy, kde, shape-validation
covariance does not have dimension {self.d}
exception error jax, scipy, kde, covariance, shape-validation
integrate_box_1d() only handles 1D pdfs
exception error jax, scipy, kde, dimensionality
the limits of integration in integrate_box_1d must be scalar
exception error jax, scipy, kde, scalar-validation
KDEs are not the same dimensionality
exception error jax, scipy, kde, dimensionality, shape-validation
only 1D box integrations are supported; use `integrate_box_1
exception error jax, scipy, kde, not-implemented, api-parity
dynamically changing the bandwidth method is not supported
exception error jax, scipy, kde, not-implemented, api-parity
gaussian_kde does not support complex coordinates
exception error jax, scipy, kde, complex-dtype, not-implemented
points have dimension {}, dataset has dimension {}
exception error jax, scipy, kde, shape-validation, dimension-mismatch
points and xi must have same trailing dim
exception error jax, scipy, kde, shape-validation
precision matrix must match data dims
exception error jax, scipy, kde, precision-matrix, shape-validation
x and n must be of integer type; got x.dtype={x.dtype}, n.dt
exception error jax, scipy, multinomial, dtype, statistics
allow_singular argument of multivariate_normal.logpdf
exception error jax, scipy, multivariate-normal, not-implemented, covariance
multivariate_normal.logpdf got incompatible shapes
exception error jax, scipy, multivariate-normal, shape-mismatch
smap axis_name should be a `str` or a `Hashable`, but got {a
exception error jax, shard-map, smap, axis-name, distributed
smap in_axes must be an int, None, jax.sharding.Infer, or a
exception error jax, smap, shard-map, in-axes, distributed
smap in_axes must be an int, None, jax.sharding.Infer, or (n
exception error jax, smap, in-axes, pytree, validation
smap out_axes must be an int, None, or (nested) container wi
exception error jax, smap, out-axes, validation
shard_map requires a callable for its first argument, but go
exception error jax, shard-map, callable, type-error
in_specs passed to shard_map: {s} does not match the specs o
exception error jax, shard-map, sharding, partition-spec, distributed
shard_map out_specs rank error (msg from _spec_rank_error)
exception error jax, shard-map, out-specs, rank, distributed
shard_map out_specs vma error (msg from _inout_vma_error)
exception error jax, shard-map, out-specs, vma, replication
The context mesh cannot be empty. Use `jax.set_mesh(mesh)` t
exception error jax, shard-map, mesh, context, distributed
The context mesh {ctx_mesh} should match the mesh passed to
exception error jax, shard-map, mesh, context-mismatch, distributed
shard_map requires a `jax.sharding.Mesh` or a `jax.sharding.
exception error jax, shard-map, mesh, type-error, distributed
shard_map requires a non-empty mesh. Got {mesh}
exception error jax, shard-map, mesh, empty-mesh, distributed
`axis_names` argument of shard_map should be of type `frozen
exception error jax, shard-map, axis-names, type-error
jax.shard_map requires axis_names={axis_names} to be a subse
exception error jax, shard-map, axis-names, mesh, distributed
in_axes was not specified when axis_name={axis_names} was of
exception error jax, shard-map, smap, axis-types, in-specs
Invalid spec: {spec}
exception error jax, shard-map, partition-spec, manual-mode
unreduced/reduced can only be passed to {prefix}_specs when
exception error jax, shard-map, manual-mode, unreduced, experimental
unreduced in {prefix}_specs {s} can only be used when the me
exception error shard-map, partition-spec, mesh, jax, unreduced
reduced in {prefix}_specs {s} can only be used when the mesh
exception error shard-map, partition-spec, mesh, jax, reduced
shard_map in_specs argument must be a pytree of `jax.shardin
exception error shard-map, partition-spec, none-default, jax, typeerror
shard_map {prefix}_specs argument must refer to an axis mark
exception error shard-map, manual-axes, partition-spec, jax
shard_map {prefix}_specs argument must be a pytree of `jax.s
exception error shard-map, partition-spec, typeerror, jax
shard_map in_specs rank error (msg from _spec_rank_error)
exception error shard-map, rank-mismatch, partition-spec, jax
shard_map in_specs divisibility error (msg from _spec_divisi
exception error shard-map, divisibility, shape-mismatch, jax
in_specs containing unreduced {spec} passed to shard_map sho
exception error shard-map, unreduced, sharding-mismatch, jax
in_specs containing unreduced_kind {spec} passed to shard_ma
exception error shard-map, unreduced-kind, sharding-mismatch, jax
in_specs containing reduced {spec} passed to shard_map shoul
exception error shard-map, reduced, sharding-mismatch, jax
out_specs passed to shard_map should be equal to the unreduc
exception error shard-map, unreduced, out-specs, sharding-mismatch, jax
out_specs passed to shard_map should be equal to the unreduc
exception error shard-map, unreduced-kind, out-specs, sharding-mismatch, jax
out_specs passed to shard_map should be equal to the reduced
exception error shard-map, reduced, out-specs, sharding-mismatch, jax
Mesh shape of the input {a.sharding.mesh.shape_tuple} does n
exception error shard-map, mesh-mismatch, named-sharding, jax
Please pass `jax.Array`s with a `NamedSharding` as input to
exception error shard-map, abstract-mesh, named-sharding, jax
Invalid value ({e.ty}) encountered in sharded computation.
exception error shard-map, nan, inf, debug-flags, jax
Eager shard_map + unreduced/reduced + partial manual is not
exception error shard-map, eager, not-implemented, unreduced, jax
custom_vjp symbolic_zeros support with shard_map is not impl
exception error shard-map, custom-vjp, symbolic-zeros, autodiff, jax
device_put with explicit device not allowed within shard_map
validation error shard-map, device-put, placement, jax
Eager shard_map cannot return a `jax.Ref`. Please wrap your
validation error shard-map, ref, jit-required, jax
vmap spmd_axis_name cannot appear in shard_map in_specs
validation error jax, vmap, shard-map, spmd, partition-spec
vmapped away explicit mesh axis cannot appear in shard_map i
validation error jax, vmap, shard-map, explicit-mesh
vmap spmd_axis_name cannot appear in shard_map out_specs
validation error jax, vmap, shard-map, spmd, out-specs
vmapped away explicit mesh axis cannot appear in shard_map o
validation error jax, vmap, shard-map, explicit-mesh, out-specs
out_sharding passed to top_level_all_gather cannot be {out_s
validation error jax, shard-map, all-gather, sharding-type
Input sharding mesh {aval.sharding.mesh} should be equal to
validation error jax, mesh, all-gather, shard-map
multiple dimensions cannot be all_gathered since multi_dim=F
validation error jax, all-gather, shard-map, multi-dim
top_level_all_gather doesn't allow input {aval} to be unshar
validation error jax, all-gather, shard-map, resharding
top_level_all_gather maintains `top_level_all_gather(x, ...)
validation error jax, all-gather, shard-map, partition-spec
top_level_all_gather works when all mesh axes of context mes
validation error jax, all-gather, explicit-mesh, context-mesh
{self} implies that array axis {dim} is partitioned {p} time
validation critical jax, sharding, indivisible-shape, pjit, distributed
Subclasses should implement this method.
exception error jax, sharding, not-implemented, abstract-method
Subclasses should implement this method
exception error jax, sharding, not-implemented, memory-kind
Invalid dtype for `swap`. Ref dtype: {expected_out_ty.dtype}
validation error jax, dtype-mismatch, state-primitives, swap
`addupdate` must be called on `Ref` types: {ref_aval}.
validation error jax, state-primitives, type-error, addupdate
Invalid shape for `addupdate`. Ref shape: {ref_aval.shape}.
validation error jax, shape-mismatch, state-primitives, in-place-update
Invalid dtype for `addupdate`. Ref dtype: {ref_aval.dtype}.
validation error jax, dtype-mismatch, state-primitives, in-place-update
Invalid sharding for `addupdate`. Ref sharding: {ref_aval.sh
validation error jax, sharding, spmd, state-primitives
performing a set/swap operation with a differentiated value
exception error jax, autodiff, state-primitives, jvp
No support for vmapping over nontrivial slices just yet: {id
exception error jax, vmap, batching, not-implemented, slicing
Batching with multiple indexers not supported.
exception error jax, vmap, batching, not-implemented, multiple-indexing
Permutation {self.permutation} does not match the rank of th
validation error jax, transpose, shape-mismatch
Cannot transpose {x} to {self.permutation}
validation error jax, internal, type-dispatch
Cannot select {ref}
validation error jax, select, type-error
Cannot select from Refs of different types: {types}
validation error jax, multiref, type-mismatch
Cannot resolve attribute {name} from: {attrs}
exception error jax, multiref, attribute-mismatch
Multi-ref TransformedRef requires a single transform: {self}
exception error jax, multiref, api-restriction
{self!r} has no `shape`.
exception error jax, attribute-error, shape
{self!r} has no `dtype`.
exception error jax, attribute-error, dtype
Bitcast ref with dynamic size is not supported.
exception error jax, bitcast, dynamic-shape
Reshape ref with dynamic size is not supported.
exception error jax, reshape, dynamic-shape
Transpose with multiref is not supported.
exception error jax, transpose, multiref
Found inconsistent memory spaces in multiref: {self.ref}
exception error jax, memory-space, multiref
Aval memory space {inner_aval.memory_space} does not match t
exception error jax, memory-space, aval
len() of unsized object
exception error jax, len, unsized
{self!r} has no `sharding`.
exception error jax, sharding, attribute-error
{self!r} has no `manual_axis_type`.
exception error jax, axis-type, attribute-error
Bitcast 1D ref with bitwidth change is not supported.
exception error jax, bitcast, shape-restriction
Expected input and output shapes are the same after multiply
exception error jax, bitcast, divisibility
subset_by_index must be a tuple of size 2.
exception error jax, tpu, svd, linalg, argument-validation
Got empty index range in subset_by_index.
exception error jax, tpu, svd, linalg, off-by-one
Indices in subset_by_index must be non-negative.
exception error jax, tpu, svd, linalg, negative-index
Index in subset_by_index[1] exceeds matrix size.
exception error jax, tpu, svd, linalg, bounds-check
full_matrices and subset_by_index cannot be both be set.
exception error jax, tpu, svd, linalg, mutually-exclusive-args
The SVD algorithm parameter is not implemented on TPU.
exception error jax, tpu, svd, not-implemented, backend-support
Only the POLAR (which is also DEFAULT on TPU) SVD algorithm
exception error jax, tpu, svd, mlir-lowering, not-implemented
tpu_custom_call does not support non-trivial batching.
exception error jax, tpu, vmap, custom-call, not-implemented
invalid memory space: {self}
exception error jax, tpu, memory-space, enum, invalid-value
input_memory_space_colors only supports HBM, VMEM and SMEM
exception error jax, tpu, custom-call, serialization, not-implemented
invalid flag value: {value}
exception error tpu, pallas, custom-call, config-serialization, jax
Mosaic kernels cannot be automatically partitioned. Please w
exception error tpu, pallas, sharding, shard-map, jax
Replica lowering for Mosaic kernels not implemented.
exception error tpu, pallas, communication, replication, jax
Unknown core type: {core_type}
exception error tpu, pallas, core-type, version-mismatch, mlir
The iteration bounds and dimension semantics attributes must
exception error tpu, pallas, mlir, subkernel, validation
The iteration bound corresponding to the core-parallel dimen
exception error tpu, pallas, dynamic-shape, static-bounds, jax
A single Mosaic subkernel cannot contain multiple core shard
exception error tpu, pallas, core-parallel, subkernel, validation
The iteration bound corresponding to the core-parallel dimen
exception error tpu, pallas, core-parallel, subkernel, consistency
The manually assigned {collective_id=} in {kernel_name=} con
exception error tpu, pallas, collective-id, barrier, conflict
collective_id has to be specified when using a custom barrie
exception error tpu, pallas, collective-id, barrier, missing-argument
collective_id has to be unspecified or None when not using a
exception error tpu, pallas, collective-id, barrier, invalid-argument
vmem_limit_bytes must be an int: provided with a {type(vmem_
exception error tpu, pallas, vmem, type-error, config
explicit tiling is only supported for SparseCore kernels.
exception error tpu, pallas, tiling, sparsecore, invalid-argument
explicit opt_level is only supported for SparseCore kernels.
exception error tpu, pallas, opt-level, sparsecore, invalid-argument
TPU generation is not supported
exception error tpu, hardware-generation, version-mismatch, jax
Unsupported TPU chip version: {chip_version}
exception error tpu, chip-version, unsupported-hardware, jax
Unsupported TPU device kind: {device_kind}. If you are not r
exception error tpu, cpu-testing, abstract-mesh, device-kind, jax
Lite chips, single core chips, and dual-core chips that do n
exception error tpu, chip-config, megacore, validation, jax
Backend {name} already initialized
exception error
Mock topology must be of the form "<number-of-slices> x <num
exception error
PJRT plugin config file should contain "library_path" field.
exception error
Multiple ':' separators for option in jax_pjrt_client_create
exception error
Cannot provide both 'factory' and 'options' when registering
exception error
Unknown backend: '{platform}' requested, but no platforms th
exception error
{err_msg} (set JAX_PLATFORMS='' to automatically choose an a
exception error
out_specs_fn already specified
validation error jax, colocated-python, api-misuse, configuration
out_specs already specified
validation error jax, colocated-python, api-misuse, configuration
devices already specified
validation error jax, colocated-python, devices, api-misuse
{op} has an unsupported layout: {out_layout_attr}
exception error jax, mosaic, gpu, layout, tiled-layout
Unsupported memory space: {orig_ref_ty.memory_space}
exception error jax, mosaic, gpu, memory-space, memref
Unsupported memory space: {ref_type.memory_space}
exception error jax, mosaic, gpu, memory-space, store
Expected TiledLayout, got {type(layout)}
exception error jax, mosaic, gpu, smem, layout
Transfer of {total_bits} bits is not divisible by {8 * utils
exception error jax, mosaic, gpu, smem, alignment, tma
`strides` must contain only 1s.
exception error jax, mosaic, gpu, vector, strides
Only slicing with static indices allowed.
exception error jax, mosaic, gpu, vector, static-index
Unsupported reduction kind: {op.kind}
exception error jax, mosaic, gpu, reduction, vector
Output layout {out_layout} must match the accumulator layout
exception error jax, mosaic, gpu, reduction, layout
Only 1 reduction dimension is supported.
exception error jax, mosaic, gpu, reduction
Unsupported layout: {src.layout}
exception error jax, mosaic, gpu, reduction, layout
Found multiple SwizzleTransformAttr
exception error jax, mosaic, gpu, swizzle, transforms
Unsupported transform: {transform}
exception error jax, mosaic, gpu, transforms, version-skew
Offsets {offsets} have lower rank than tiling {tiling}
exception error jax, mosaic, gpu, tiling, offsets, rank-mismatch
Offset {i} is not divisible by tile size {t}
exception error jax, mosaic, gpu, tiling, alignment
Strides {strides} have lower rank than tiling {tiling}
exception error jax, mosaic, gpu, tiling, strides, rank-mismatch
Can not tile strides when tiled dimensions have been transpo
exception error jax, mosaic, gpu, tiling, strides, transpose
Stride {s} is not divisible by {d} (tile size = {t}). Stride
exception error jax, mosaic, gpu, tiling, strides, alignment
SubViewOp only supports a single tile transform.
exception error jax, mosaic-gpu, memref, subview, tile-transform
Only support memref.cast where the input and output types ar
exception error jax, mosaic-gpu, memref, cast, type-mismatch
memref.cast transforms must have identical transforms for bo
exception error jax, mosaic-gpu, memref, cast, shared-memory, transforms
memref.cast tmem layouts must be identical for both input an
exception error jax, mosaic-gpu, tmem, cast, layout
Unsupported memory space when lowering memref.cast: {memory_
exception error jax, mosaic-gpu, memref, cast, memory-space
Swizzle mismatch. In transforms swizzle: {in_swizzle}, out t
validation error jax, mosaic-gpu, transpose, swizzle, transforms
Size mismatch for in/out transforms. In transforms: {in_tran
validation error jax, mosaic-gpu, transpose, transforms, count-mismatch
Unsupported in/out transforms. In transform: {in_transform},
exception error jax, mosaic-gpu, transpose, tile-transform, not-implemented
Cannot tile a transpose ({permutation}). Tiling dims ({permu
validation error jax, mosaic-gpu, transpose, tiling, permutation
Invalid in/out transforms. In transform: {in_transform}, out
validation error jax, mosaic-gpu, transpose, tiling, consistency
Expanding tiled dimensions is not supported.
exception error jax, mosaic-gpu, expand-shape, tiling, not-implemented
Expected the same number of in/out transforms, but got {in_t
validation error jax, mosaic-gpu, collapse-shape, transforms, count-mismatch
Only a single tiling transform is supported when collapsing
exception error jax, mosaic-gpu, collapse-shape, tile-transform, not-implemented
Collapsing the shape of a memref with non-contiguous strides
exception error jax, mosaic-gpu, collapse-shape, strides, contiguity
Input tiling {t_in.tiling} is not compatible with {op}
validation error jax, mosaic-gpu, collapse-shape, tiling, unsatisfiable
Input/output tiling mismatch when attempting to collapse a s
validation error jax, mosaic-gpu, collapse-shape, tiling, consistency
Reassociation {reassociation} is not compatible with tiling
validation error jax, mosaic-gpu, collapse-shape, reassociation, tiling
memref.LoadOp does not support transforms: {op}
exception error jax, mosaic-gpu, memref-load, transforms, not-implemented
memref.StoreOp does not support transforms: {op}
exception error jax, mosaic-gpu, memref-store, transforms, not-implemented
{ref} is not a memref.
validation error mosaic, gpu, tmem, mlir, type-validation
{ref} has a memory space {mem_ref_ty.memory_space} that is n
validation error mosaic, gpu, tmem, memory-space, mlir
{ref} has a layout {layout_attr} that does not match the exp
validation error mosaic, gpu, tmem, layout-mismatch
Expected both or neither of scales to be specified.
validation error mosaic, gpu, tcgen05, mma, scales, argument-validation
A custom return op must terminate the block.
validation error mosaic, gpu, custom-primitive, mlir, terminator
Unsupported op type: {op}
exception error mosaic, gpu, scf, control-flow, not-implemented
Input layouts {in_layouts} do not match yield layouts {yield
validation error mosaic, gpu, scf, while-loop, layout-mismatch
Output layouts {out_layouts} do not match condition layouts
validation error mosaic, gpu, scf, while-loop, layout-mismatch
Rule cannot handle an op with vector operands or results: {o
validation error mosaic, gpu, lowering, layout-mismatch, not-implemented
gpu.launch op not found.
validation error mosaic, gpu, gpu-launch, pass-pipeline, mlir
Tiles must have a decreasing rank
validation error mosaic, fragmented-array, tiling, layout, validation
Tiles must not be empty
validation error mosaic, fragmented-array, tiling, validation
Tile shape must only have positive sizes, got: {self.tiles}
validation error mosaic, fragmented-array, tiling, validation
Tiling {self.tiles} does not apply to shape {orig_shape}
validation error mosaic, fragmented-array, tiling, shape-mismatch
shape {orig_shape} is not a valid result of applying tiling
validation error mosaic, fragmented-array, tiling, shape-mismatch
Invalid dimension {dim} for tiling {self}
validation error mosaic, fragmented-array, tiling, index-out-of-range
Shape {shape} and strides {strides} must have the same lengt
exception error mosaic, fragmented-array, tiling, shape-strides, validation
`dimension` must be smaller than the rank of the array.
exception error jax, mosaic-gpu, iota, dimension-validation, shape-mismatch
Unsupported bitwidth: {bitwidth}
exception error jax, mosaic-gpu, layout-cast, bitwidth, wgmma, unsupported-dtype
Cannot convert from {self.layout} to {new_layout}
exception error jax, mosaic-gpu, layout-conversion, not-implemented, wgmma
Pointwise operations on {bitwidth}-bit types are unsupported
exception error jax, mosaic-gpu, pointwise, quantization, int4, int8, dtype-bitwidth
Cannot broadcast shape {self.shape} to layout {o.layout}
exception error jax, mosaic-gpu, broadcasting, splat-layout, shape-mismatch
Incompatible FragmentedArray layouts
exception error jax, mosaic-gpu, layout-mismatch, pointwise, fragmented-array
Incompatible FragmentedArray shapes
exception error jax, mosaic-gpu, shape-mismatch, registers, pointwise
trunc_div supports only integer types, got {self.mlir_dtype}
validation error jax, mosaic-gpu, division, dtype-check, integer-only
{arg_ty}
exception error mosaic, gpu, approx, internal
{vec_bitwidth}
exception error mosaic, gpu, ptx, packed-instruction
output_is_signed must be non-None if and only if the MLIR ty
validation error mosaic, gpu, bitcast, type-error, signedness
Only bitcast between types of the same bitwidth supported
validation error mosaic, gpu, bitcast, bitwidth
Only arrays with tiled layouts can be sliced
exception error mosaic, gpu, slicing, layout
Only slicing with static indices allowed
validation error mosaic, gpu, slicing, static-indices
Integer indexing not implemented for tiled dimensions (only
exception error mosaic, gpu, slicing, tiled-layout
Base indices of array slices must be aligned to the beginnin
validation error mosaic, gpu, slicing, tile-alignment
The slice shape must be a multiple of the tile shape. The ar
validation error mosaic, gpu, slicing, tile-alignment
Source and destination layouts aren't compatible for a broad
exception error jax, mosaic-gpu, layout, broadcast, fragmented-array
Splat layout does not support multimem
exception error jax, mosaic-gpu, splat-layout, multimem, tma
Atomic stores not supported for splat layout
exception error jax, mosaic-gpu, atomic, splat-layout, store
Only TiledLayouts support swizzling
exception error jax, mosaic-gpu, swizzle, strided-layout, store
Arrays with the splat layout can only be stored when they ha
exception error jax, mosaic-gpu, splat-layout, warpgroup, alignment, store
Multimem refs are not supported in store_tiled_async
exception error jax, mosaic, gpu, multimem, async-store, not-implemented
Replicated dimensions are not supported
exception error jax, mosaic, gpu, layout, replicated, warp, not-implemented
Unsupported register bitwidth: {reg_bitwidth}
exception error jax, mosaic, gpu, ptx, register-width, not-implemented
f32 not supported for async atomics
exception error jax, mosaic, gpu, atomics, f32, cluster-barrier, not-implemented
f32 only supports add atomics, got {atomic}
exception error jax, mosaic, gpu, atomics, f32, not-implemented
f16/bf16 not supported for async atomics
exception error jax, mosaic, gpu, atomics, f16, bf16, cluster-barrier, not-implemented
f16/bf16 only supports add, min, max atomics, got {atomic}
exception error jax, mosaic, gpu, atomics, f16, bf16, not-implemented
f16/bf16 SMEM/multimem atomics only support add, got {atomic
exception error jax, mosaic, gpu, atomics, f16, bf16, shared-memory, multimem, not-implemented
Unsupported element type for atomic stores: {element_type}
exception error jax, mosaic, gpu, atomics, dtype, not-implemented
Unsupported copy: {src.type} -> {dst.type}
exception error jax, mosaic-gpu, gpu-kernel, shared-memory, not-implemented
Need at least one array to concatenate
validation error jax, mosaic-gpu, validation, empty-input
{axis=} is out of bounds for array of {rank=}
validation error jax, mosaic-gpu, validation, axis-bounds
All arrays must have the same rank, got {len(arr.shape)} at
validation error jax, mosaic-gpu, validation, rank-mismatch
All arrays must have the same dtype, got {arr.mlir_dtype} at
validation error jax, mosaic-gpu, validation, dtype-mismatch
All arrays must have the same signedness, got {arr.is_signed
validation error jax, mosaic-gpu, validation, signedness
All arrays must have matching shapes along non-concatenated
validation error jax, mosaic-gpu, validation, shape-mismatch
All arrays must have the same layout, got {arr.layout} at in
validation error jax, mosaic-gpu, validation, layout-mismatch
Concatenating arrays with strided layout is only supported a
exception error jax, mosaic-gpu, not-implemented, wgmma, layout-limitation
Expected WGStridedFragLayout, got {arr.layout} at index {i}
validation error jax, mosaic-gpu, validation, layout-mismatch, wgmma
All WGStridedFragLayout arrays must have the same vec_size,
validation error mosaic-gpu, fragmented-array, layout-mismatch, concat, validation
Concatenating arrays with splat layout is not supported.
exception error mosaic-gpu, not-implemented, splat-layout, concat, fragmented-array
{op} does not have an {name} attribute
exception error mosaic-gpu, mlir, missing-attribute, layout-inference
{op} has {name} of an unexpected type: {result}
exception error mosaic-gpu, mlir, type-mismatch, attribute-type, layout-inference
Unknown attribute: {attr_name}
exception error mosaic-gpu, invalid-argument, attribute-dispatch, internal-api
Subclasses should override this method
exception error mosaic-gpu, not-implemented, abstract-method, memref-transform, subclassing
The collective size ({collective_size}) must divide the slic
validation error jax, mosaic-gpu, async-copy, cluster-partitioning, shape-validation
async_copy requires all GMEM strides except the last one to
validation error jax, mosaic-gpu, tma, memory-alignment, strides
Async copies only support striding up to 5 dimensions
validation error jax, mosaic-gpu, tma, rank-limit, shape-validation
None of the leading dimensions in the transformed slice shap
validation error jax, mosaic-gpu, tma, cluster-partitioning, shape-validation
Async copies require the number of bits copied along the las
validation error jax, mosaic-gpu, tma, memory-alignment, shape-validation
Async copies with {swizzle=} require the last dimension of t
validation error jax, mosaic-gpu, tma, swizzling, shared-memory
Expected same element type, got {element_type} and {dst_ref_
validation error jax, mosaic-gpu, async-copy, dtype-mismatch
Only the TMA implementation supports reductions
validation error jax, mosaic-gpu, tma, reduction, unsupported-operation
Reduction op {reduction_op} not supported by the TMA impleme
validation error jax, mosaic-gpu, tma, reduction, dtype-mismatch, unsupported-operation
Barriers are required for TMA GMEM -> SMEM copies
validation error jax, mosaic-gpu, tma, barrier, synchronization
Only unit strides are supported but got {op.static_strides}.
validation error jax, mosaic-gpu, memref, layout-inference, strides, not-implemented
Transposed memrefs are not supported in ExpandShapeOp.
validation error jax, mosaic-gpu, memref, expand-shape, transpose, layout-inference
CollapseShapeOp with empty reassociation is not supported.
validation error jax, mosaic-gpu, memref, collapse-shape, layout-inference, not-implemented
CollapseShapeOp with non-contiguous strides is not supported
validation error jax, mosaic-gpu, memref, collapse-shape, strides, layout-inference
Only scalar memrefs are supported, got {ref_shape}
validation error jax, mosaic-gpu, memref, load-store, layout-inference, scalar
Unsupported transforms {transforms}
validation error jax, mosaic-gpu, smem-transforms, swizzle, tiling, layout-inference
Cannot apply {tiling_transform_str} with {swizzle_str} to me
validation error jax, mosaic-gpu, smem-transforms, swizzle, tiling, shape-mismatch
Unsupported TMA index shape {shape}
validation error jax, mosaic-gpu, tma, async-load, indices, gather
Only leading gather dimensions allowed.
validation error jax, mosaic-gpu, tma, gather, async-load, layout-inference
Only 2D gathers/scatters for async load/store are supported.
validation error jax, mosaic-gpu, tma, gather, async-load, rank, layout-inference
Cannot assign layout to async load with gather indices since
validation error jax, mosaic, gpu, layout, alignment, async-load
Expected the same number of in_{attr_suffix} ({len(in_layout
validation error jax, mosaic, mlir, layout, invariant
Expected the same number of out_{attr_suffix} ({len(out_layo
validation error jax, mosaic, mlir, layout, invariant
Inferred layout not found for operand {operand}.
exception error jax, mosaic, layout-inference, internal
Inferred layout {operand_layout} for operand {operand} does
exception error jax, mosaic, layout, conflict
No layout inference rule defined for {op}
exception error jax, mosaic, not-implemented, unsupported-op
Failed to infer a possible set of layouts. This should only
exception error jax, mosaic, layout, unsatisfiable-constraints
Unsupported layout for conversion from MLIR attribute: {attr
exception error jax, mosaic, layout, mlir-attribute, not-implemented
Unsupported transform {transform}
exception error jax, mosaic, transform, not-implemented
m_warps must be 1, 2, or 4, but got {m_warps=}
exception error jax, mosaic, mma, warps, validation
is_signed must be specified for integer types
exception error jax, mosaic, mma, dtype, signedness
M mismatch: {m} != {m2}
exception error jax, mosaic, mma, shape-mismatch
N mismatch: {n} != {n2}
exception error jax, mosaic, mma, shape-mismatch
K mismatch: {k} != {k2}
exception error jax, mosaic, mma, shape-mismatch
Dtype mismatch: {a.mlir_dtype} != {b.mlir_dtype}
exception error jax, mosaic, mma, dtype-mismatch
Unsupported operand type: {element_type}
exception error jax, mosaic, mma, unsupported-dtype
Only s32 accumulator supported for integer operands.
exception error jax, mosaic, mma, accumulator, int8
Only signed accumulator supported for integer operands.
exception error jax, mosaic, mma, accumulator, signedness
{iterations=} must be positive
validation error profiling, mosaic, gpu, argument-validation
Unsupported trace scope: {trace_scope}
validation error profiling, mosaic, gpu, argument-validation
Scope {self.trace_scope} not supported
exception error profiling, mosaic, gpu, immutability
Block size is not a multiple of {scope_size}
validation error profiling, mosaic, gpu, cuda-launch-config
Allocated too many names
exception error profiling, mosaic, gpu, resource-exhaustion
Insufficient space to capture a full trace
exception error profiling, mosaic, gpu, buffer-overflow
Profiler collected an odd number of trace events. This likel
exception error profiling, mosaic, gpu, memory-corruption
Scope {spec.trace_scope} not supported
exception error profiling, mosaic, gpu, argument-validation
Scope {self.spec.trace_scope} not supported
exception error profiling, mosaic, gpu, immutability
Unsupported accumulator dtype: {acc_dtype}
exception error tcgen05, matmul, mosaic, gpu, dtype
Unsupported input dtype: {ty}
exception error tcgen05, matmul, mosaic, gpu, dtype
N must be a multiple of 8 and <= 256, got: {n}
validation error tcgen05, matmul, mosaic, gpu, shape-validation
M must be a multiple of 16 and <= 256, got: {m}
validation error tcgen05, matmul, mosaic, gpu, shape-validation
Unsupported scale type: {scale_type}
exception error tcgen05, mx-format, matmul, dtype
No swizzle is not supported
exception error tcgen05, matmul, swizzle, shared-memory
Either none or both scales should be provided
validation error tcgen05, matmul, mx-format, argument-validation
Sparse MMA not supported for M=64
exception error jax, mosaic, gpu, sparse, mma, tensor-core, tcgen05, blackwell
Swizzle={b_swizzle} is too big for MMA with M=64. Try loweri
validation error jax, mosaic, gpu, mma, swizzle, tile-size, tcgen05
Only M=128 and M=64 are supported for MMA, but got M={m}
validation error jax, mosaic, gpu, mma, shape-validation, tcgen05
Sparse MMA unsupported for f32
exception error jax, mosaic, gpu, sparse, dtype, mma, tcgen05
MMA with element type {elem_type_str} does not support block
validation error jax, mosaic, gpu, block-scaling, dtype, mma, tcgen05, mxfp
MMA with element type {elem_type_str} only supports accumula
validation error jax, mosaic, gpu, accumulator-dtype, dtype, mma, tcgen05
MMA with element type {elem_type_str} only supports accumula
validation error jax, mosaic, gpu, accumulator-dtype, dtype, mma, tcgen05
Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn,
validation error gpu, mosaic, tcgen05, dtype, block-scaling
Unsupported element type for block scaling: {a_element_type}
validation error gpu, mosaic, tcgen05, block-scaling, not-implemented
A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}),
validation error gpu, mosaic, tcgen05, shape-mismatch, block-scaling
A scale layout {a_scale.layout} is not supported
validation error gpu, mosaic, tcgen05, layout, tmem
Expected B scales to have a M=64 collective layout, got {b_s
validation error gpu, mosaic, tcgen05, layout, collective, cgmma
Expected B scales to have a M=128 layout, got {b_scale.layou
validation error gpu, mosaic, tcgen05, layout, block-scaling
B scale shape[0] must be a multiple of 128 and >= N={n * num
validation error gpu, mosaic, tcgen05, shape-mismatch, block-scaling, alignment
B scale shape mismatch: expected ({b_scale.shape[0]}, {k_sca
validation error gpu, mosaic, tcgen05, shape-mismatch, block-scaling
A sparse metadata shape mismatch: expected {(m, expected_met
validation error gpu, mosaic, tcgen05, sparsity, shape-mismatch, metadata
A sparse metadata dtype mismatch: expected i2, got {a_sparse
validation error gpu, mosaic, tcgen05, sparsity, dtype, metadata
4-bit block scaled MMA only supports K-fastest operands, but
validation error gpu, mosaic, tcgen05, layout, mxfp4, block-scaling
4-bit block scaled MMA only supports K-fastest operands, but
validation error gpu, mosaic, tcgen05, layout, mxfp4, block-scaling
B tiling too small. Increase swizzle or transpose the input.
validation error gpu, mosaic, tcgen05, sparsity, swizzle, not-implemented
A address calculation for multiple M tiles
validation error gpu, mosaic, tcgen05, tmem, not-implemented, mma
A sparse metadata address calculation for multiple tiles
validation error gpu, mosaic, tcgen05, sparsity, not-implemented, tiling
A scale address calculation for multiple M tiles
validation error gpu, mosaic, tcgen05, block-scaling, not-implemented, tiling
B scale address calculation for multiple N tiles
validation error gpu, mosaic, tcgen05, block-scaling, not-implemented, tiling
D address calculation for multiple M tiles
exception error gpu, mosaic, tcgen05, tmem, not-implemented, tiling
Minor dimension of shape must be divisible by packing, got:
exception error jax, mosaic, tmem, layout, shape-validation
Unsupported shape: {shape}. TMEM references must have either
exception error jax, mosaic, tmem, shape-validation, blackwell
Packing must be a power of 2, got: {packing}
exception error jax, mosaic, tmem, layout, argument-validation
Columns must be a multiple of 16, got: {columns}
exception error jax, mosaic, tmem, layout, alignment
tmem_addr_ref must be a memref or a pointer, got: {tmem_addr
exception error jax, mosaic, tmem, memref, argument-validation
tmem_addr_ref must be in shared memory, got: {addr_ref_ty}
exception error jax, mosaic, tmem, shared-memory, address-space
tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}
exception error jax, mosaic, tmem, dtype, memref
tmem_addr_ref must contain a single element, got: {addr_ref_
exception error jax, mosaic, tmem, memref, shape-validation
TMEM refs must have at least 32 rows, got: {shape[0]}
exception error jax, mosaic, tmem, shape-validation, hardware-constraint
collective argument must be provided when TMEM layout is inf
exception error jax, mosaic, tmem, missing-argument
TMEM can only be sliced, not indexed
exception error jax, mosaic, tmem, slicing, unsupported-operation
TMEM cannot be sliced along rows
exception error jax, mosaic, tmem, slicing, unsupported-operation
Cannot slice TMEM with multiple tiles along rows.
exception error jax, mosaic, tmem, slicing, layout, unsupported-operation
TMEM layout {self.layout} is not supported
exception error jax, mosaic, tmem, load, layout, unsupported-operation
Unsupported reduction for i32. Only min and max are supporte
exception error jax, mosaic, tmem, reduction, dtype, unsupported-operation
Unsupported reduction for f32. Only min, max, absmin, and ab
exception error mosaic, gpu, tcgen05, reduction, tensor-core, jax
Unsupported dtype for reduction: {self.dtype}
exception error mosaic, gpu, tcgen05, dtype, reduction, jax
Sparse meta layout loads unsupported.
exception error mosaic, gpu, tcgen05, sparse, tensor-memory, not-implemented, jax
Loading multiple row tiles
exception error mosaic, gpu, tcgen05, layout, tensor-memory, not-implemented, jax
Fused load-reduce is not supported for this layout
exception error mosaic, gpu, tcgen05, reduction, layout, jax
Loads from TMEM layout {self.layout} to register layout {lay
exception error mosaic, gpu, tcgen05, layout, tensor-memory, jax
TMEM stores expect a FragmentedArray, got: {value}
exception error mosaic, gpu, tcgen05, type-check, tensor-memory, jax
Stored array has shape {value.shape}, but TMEM has shape {se
exception error mosaic, gpu, tcgen05, shape-mismatch, tensor-memory, jax
Stored array has dtype {value.mlir_dtype}, but TMEM has dtyp
exception error mosaic, gpu, tcgen05, dtype-mismatch, tensor-memory, jax
Stored array has layout {value.layout}, but TMEM stores expe
exception error mosaic, gpu, tcgen05, layout, type-check, tensor-memory, jax
Sparse meta layout stores unsupported.
exception error mosaic, gpu, tcgen05, sparse, not-implemented, tensor-memory, jax
Storing from register layout {value.layout} to TMEM layout {
exception error mosaic, gpu, tcgen05, layout, tensor-memory, jax
Unsupported packing: {self.packing}
exception error mosaic, gpu, tcgen05, debug, packing, not-implemented, jax
Unsupported packing: {tmem_packing} for element type {elt_bi
exception error mosaic, gpu, tcgen05, packing, not-implemented, jax
Only {reg_packing} packing supported for bitwidth {elt_bitwi
exception error mosaic, gpu, tcgen05, packing, not-implemented, jax
Only {reg_packing} supported for element type {dtype}, but g
exception error mosaic, gpu, tcgen05, packing, not-implemented, jax
Unsupported reduction kind: {reduce}
exception error mosaic, gpu, tcgen05, reduction, internal, jax
K tile stride must be a multiple of 16
validation error gpu, mosaic, tcgen05, sparse, shared-memory, alignment
Unsupported swizzle, expected 16, 32, 64 or 128, but got: {s
validation error gpu, mosaic, tcgen05, swizzle, validation
tcgen05.cp only supports fully packed TMEM references (packi
validation error gpu, mosaic, tcgen05, packing, tensor-memory
TMEM reference must have {TMEM_ROWS} rows, but got {tmem_ref
validation error gpu, mosaic, tcgen05, tensor-memory, shape-mismatch
Only standard TMEM layout is supported, got: {tmem_ref.layou
validation error gpu, mosaic, tcgen05, layout, tensor-memory
SMEM has shape {smem_shape}, but expected {expected_smem_sha
validation error gpu, mosaic, tcgen05, shape-mismatch, swizzle, shared-memory
The SMEM tiles must be contiguous
validation error gpu, mosaic, tcgen05, strides, shared-memory, contiguity
name must be non-empty
validation warning mosaic, gpu, io, dump, argument-validation
address_space not supported: {address_space}
exception error gpu, mosaic, address-space, nvptx, lowering
Pointer address space {ptr_ty.address_space} does not match
validation error gpu, mosaic, address-space, pointer, memref, llvm
Non-zero offset is not supported for ptr_as_memref
validation error gpu, mosaic, memref, offset, llvm
Empty array
validation error gpu, mosaic, tma, empty-array, argument-validation
Can't print the type {arg.type}
exception error gpu, mosaic, debug-print, unsupported-type, dtype
Cannot specify scope to a non-uniform debug_print.
validation error gpu, mosaic, debug-print, argument-validation
2D+ vectors are not supported in debug_print: {vec_ty}
exception error gpu, mosaic, debug-print, vector, rank
Only 32-, 64- and 128-bit stores are supported
validation error gpu, mosaic, dsmem, multimem, store, bit-width
Folding {fold_rank} dimensions starting from {dim} is out of
validation error jax, mosaic-gpu, memref, shape-mismatch, index-out-of-bounds
strides={ref_ty.get_strides_and_offset()[0]}, {ref_ty.shape=
validation error jax, mosaic-gpu, memref, strides, layout
Can only infer one dimension
validation error jax, mosaic-gpu, memref, reshape, invalid-argument
Non-divisible unfold:
validation error jax, mosaic-gpu, memref, reshape, non-divisible-shape
Index {idx} along axis {axis} is out of bounds for shape {sh
validation error jax, mosaic-gpu, memref, index-out-of-bounds, indexerror
Strided slices not implemented
validation error jax, mosaic-gpu, memref, slice, not-implemented, strided
Slice {idx} along axis {axis} is out of bounds for shape {sh
validation error jax, mosaic-gpu, memref, slice, index-out-of-bounds
Expected an index-typed index
validation error jax, mosaic-gpu, memref, mlir, type-mismatch
Only up to 32 barriers per group supported
validation error jax, mosaic-gpu, barrier, hardware-limit, cuda
Barrier offset {offset} is out of bounds
validation error jax, mosaic-gpu, barrier, index-out-of-bounds
Expected a dynamic index or an integer, got {offset}
validation error jax, mosaic-gpu, barrier, type-mismatch
Unsupported scope: {scope}
validation error jax, mosaic-gpu, barrier, scope, invalid-argument
Can only await on CTA-local barriers
validation error jax, mosaic-gpu, barrier, scope, ptx
Only single-thread arrival is supported on pre-Hopper hardwa
validation error jax, mosaic-gpu, barrier, gpu-architecture, hopper
Predicate not supported for no-complete arrive
validation error jax, mosaic-gpu, barrier, predicate, not-implemented
arrive_expect_tx is only supported on Hopper+ hardware
validation error jax, mosaic-gpu, barrier, tma, gpu-architecture, hopper
complete_tx is only supported on Hopper+ hardware
validation error jax, mosaic-gpu, barrier, tma, gpu-architecture, hopper
Expected tiling to be at most rank of shape. Got tiling: {ti
validation error mosaic-gpu, tiling, shape-validation
Non-divisible tiling:
validation error mosaic-gpu, tiling, divisibility
Warp reduction group size should be a power of 2 (got {group
validation error mosaic-gpu, warp-reduction, power-of-two
{memref_ty} {static_offset=} is not divisible by {packing=}`
validation error mosaic-gpu, memref, alignment, sub-byte-types
Only 32-bit types supported
validation error mosaic-gpu, redux, dtype, hardware-limit
Only 32-bit scalar types supported
validation error mosaic-gpu, redux, dtype, hardware-limit
F32 redux only supported on Blackwell GPUs
validation error mosaic-gpu, redux, gpu-architecture, blackwell
Types must match, got {high.type} and {low.type}
validation error mosaic-gpu, prmt, type-mismatch
Can't bitcast {x.type} (of bitwidth {x_bw}) to {new_type} (o
validation error mosaic-gpu, bitcast, bitwidth
Can't bitcast {x.type} to {new_type}
validation error mosaic-gpu, bitcast, vector-types
Only 1D vectors are supported {v_ty}
validation error mosaic-gpu, vector-slice, rank, not-implemented
Cannot concatenate an empty list of vectors
validation error mosaic-gpu, vector-concat, empty-input
Cannot concatenate non-vector values
validation error mosaic-gpu, vector-concat, type-validation
Cannot concatenate vectors of different element types
validation error mosaic-gpu, vector-concat, dtype-mismatch
Can only concatenate 1D vectors
validation error mosaic-gpu, vector-concat, rank
Expected a memref type but got {ref}
validation error mosaic-gpu, memref, type-validation, shared-memory
Expected a single dimension when passing a single index
validation error jax, mosaic-gpu, cluster, api-misuse
Expected SMEM but got: {ref_ty.memory_space}
validation error jax, mosaic-gpu, smem, memory-space, gpu-clusters
Cannot retrieve the architecture without an insertion point
validation error jax, mosaic-gpu, mlir, insertion-point, architecture
Cannot retrieve the architecture: no module found
validation error jax, mosaic-gpu, architecture, mlir-module
WGMMA requires m and n to be multiples of 64 and 8, got {m}
validation error jax, mosaic-gpu, wgmma, gpu-kernel, shape-validation
PTX does not support unsigned WGMMA accumulators
validation error jax, mosaic-gpu, wgmma, ptx, signedness
Only WGMMA layouts supported in WGMMAAccumulator
validation error jax, mosaic-gpu, wgmma, fragmented-array, layout
Unsupported wgmma types {(out_ty, a_element_type)=}
validation error jax, mosaic-gpu, wgmma, dtype, tensor-core
Unsupported wgmma types {(out_ty, b_element_type)=}
validation error jax, mosaic-gpu, wgmma, dtype, tensor-core
Only f16 WGMMA supports transposes
validation error jax, mosaic-gpu, wgmma, transpose, dtype
Unsupported A register array dtype: {a.mlir_dtype}
validation error jax, mosaic-gpu, wgmma, dtype, registers
Unsupported A register array shape
validation error jax, mosaic-gpu, wgmma, shape, swizzle
Unsupported A register array layout
validation error jax, mosaic-gpu, wgmma, layout, fragmented-array
Unsupported WGMMA features with A in registers
validation error jax, mosaic-gpu, wgmma, api-misuse, registers
WGMMA instruction only supports f32, f16 and s32 out (got {o
validation error jax, mosaic-gpu, wgmma, dtype, accumulator
No swizzle is not supported
validation error jax, mosaic-gpu, wgmma, swizzle, smem
Existing output dtype must match preferred_element_type.
validation error jax, pallas, tpu, megablox, dtype-mismatch
group_offset must be a ()-shaped array. Got: {group_offset.s
validation error jax, pallas, tpu, megablox, shape-validation
No tuned tiling found for (m, k, n) = ({m}, {k}, {n})
validation error jax, pallas, tpu, megablox, tiling, unsupported-shape
k_pages and v_pages must have the same shape. Got {k_pages.s
validation error jax, pallas, tpu, paged-attention, shape-validation
Number of Q heads must be divisible by number of KV heads. G
validation error jax, pallas, tpu, paged-attention, gqa, head-dim-validation
head_dim of Q must be the same as that of K/V. Got {head_dim
validation error jax, pallas, tpu, paged-attention, head-dim-validation
pages_per_compute_block must be divisible by pages per seque
validation error jax, pallas, tpu, paged-attention, layout-validation
`lengths` and `q` must have the same batch size
validation error jax, pallas, tpu, paged-attention, shape-validation
`page_indices` and `q` must have the same batch size
validation error jax, pallas, tpu, paged-attention, shape-validation
The dtype of `lengths` must be int32. Got {lengths.dtype}
validation error jax, pallas, tpu, paged-attention, dtype-validation
number of KV heads must be even when megacore_mode is 'kv_he
validation error jax, pallas, tpu, paged-attention, megacore, gqa
batch size must be even when megacore_mode is 'batch'
validation error jax, pallas, tpu, paged-attention, megacore, batch-size
megacore_mode must be one of ['kv_head', 'batch', None]
validation error jax, pallas, tpu, paged-attention, megacore, invalid-argument
{num_seqs[0]=} must be less or equal to {max_num_seqs=}
validation error jax, pallas, tpu, ragged-attention, capacity-validation
{pages_per_seq=} must be greater or equal to {min_pages_per_
validation error jax, pallas, tpu, ragged-attention, kv-cache, capacity-validation
Total q tokens {cu_q_lens[num_seqs[0]]} must be less or equa
validation error jax, pallas, tpu, ragged-attention, cumulative-length-validation
{q_len=} must be less or equal to {kv_len=} at sequence {i}.
validation error jax, pallas, tpu, ragged-attention, length-validation
{num_seqs.shape=} must be (1,)
validation error jax, pallas, tpu, ragged-attention, shape-validation
Q head_dim {head_dim} must be the same as that of K/V {head_
validation error jax, pallas, tpu, ragged-attention, head-dim-validation
Expected {kv_lens.shape=} to be ({max_num_seqs},) where `max
validation error jax, pallas, tpu, paged-attention, shape-validation
Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},) whe
validation error jax, pallas, tpu, ragged-attention, shape-validation
The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must
validation error jax, pallas, tpu, dtype-validation, paged-attention
{num_q_heads=} must be divisible by {num_kv_heads=}
validation error jax, pallas, tpu, gqa, attention, head-dimension
{sliding_window=} must be positive.
validation error jax, pallas, tpu, sliding-window, config-validation
{soft_cap=} must not be 0.0.
validation error jax, pallas, tpu, softcap, attention, config-validation
{num_kv_pages_per_block=} must be in range (0, {pages_per_se
validation error jax, pallas, tpu, block-size, paged-attention
{num_queries_per_block=} must be positive.
validation error jax, pallas, tpu, block-size, validation
{vmem_limit_bytes=} must be positive.
validation error jax, pallas, tpu, memory-limit, validation
Not implemented: {num_combined_kv_heads=} can not be XLA ful
validation error jax, pallas, tpu, attention, head-tiling, not-implemented
TPU version must be 4 or higher.
exception error jax, pallas, tpu, hardware-version, not-implemented
Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).ma
validation error jax, pallas, tpu, random, philox, shape-limit
Shape dimension {shape[-2:]} must be divisible by {block_siz
validation error jax, pallas, tpu, random, philox, alignment
Offset must be scalar, got {offset.shape}
validation error jax, pallas, tpu, random, philox, scalar-validation
Only 32-bit PRNG supported.
validation error jax, pallas, tpu, random, philox, bit-width
Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).ma
validation error jax, pallas, tpu, random, threefry, shape-limit
Shape dimension {shape[-2:]} must be divisible by {block_siz
validation error jax, pallas, tpu, random, threefry, alignment
Only 32-bit PRNG supported.
validation error jax, pallas, tpu, random, threefry, bit-width
Higher-order AD not supported.
exception error jax, pallas, tpu, splash-attention, autodiff, higher-order-ad
Block sizes for dq kernel are not needed with a fused kernel
validation error jax, pallas, tpu, splash-attention, config-validation
Mask function must return a boolean-valued array, but got: {
validation error jax, pallas, tpu, splash-attention, dtype, mask
block_kv must be a multiple of {NUM_LANES}
exception error jax, pallas, tpu, splash-attention, segment-ids, block-size
block_q must be a multiple of {NUM_LANES}
exception error jax, pallas, tpu, splash-attention, segment-ids, block-size
{bkv_compute=} should be a multiple of {NUM_LANES}
exception error jax, pallas, tpu, splash-attention, block-size
partial_mask_blocks must be of type np.bool_ but got {partia
validation error jax, pallas, tpu, splash-attention, dtype, mask
Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instea
validation error jax, pallas, tpu, splash-attention, shape, rank
Expected 'key' head dimension to be: {head_dim_qk}. Instead
validation error jax, pallas, tpu, splash-attention, shape, head-dim
In MHA, expected number of 'key' heads ({num_kv_heads}) to b
validation error jax, pallas, tpu, splash-attention, gqa, shape
Expected 'key' {k.shape} and 'value' {v.shape} to have the s
validation error jax, pallas, tpu, splash-attention, shape
{bkv=} must be a multiple of {bkv_compute=}.
validation error jax, pallas, tpu, splash-attention, block-size
{bkv_compute=} must be a multiple of {NUM_LANES}.
validation error jax, pallas, tpu, splash-attention, block-size
Invalid shape for q segment_ids: {segment_ids.q.shape}. Expe
validation error jax, pallas, tpu, splash-attention, segment-ids, shape
Invalid shape for kv segment_ids: {segment_ids.kv.shape}. Ex
validation error jax, pallas, tpu, splash-attention, segment-ids, shape
Higher-order AD not supported
exception error jax, pallas, tpu, splash-attention, autodiff
{bq=} should not be greater than {q_seq_len=}
validation error jax, pallas, tpu, splash-attention, block-size, sequence-length
{bkv=} should not be greater than {kv_seq_len=}
validation error jax, pallas, tpu, splash-attention, block-size, sequence-length
{bkv=} must be a multiple of {NUM_LANES}.
validation error jax, pallas, tpu, splash-attention, block-size, backward
{bkv_compute=} should not be greater than {bkv=}
validation error jax, tpu, pallas, splash-attention, block-sizes, validation
{bkv=} should be a multiple of {bkv_compute=}
validation error jax, tpu, pallas, splash-attention, block-sizes, divisibility
Need to specify backward blocks.
validation error jax, pallas, splash-attention, block-sizes, backward
The sharding must divide the mask blocks evenly between devi
validation error jax, tpu, splash-attention, sharding, dynamic-mask
Sharding the kv sequence dimension is not supported
validation error jax, tpu, splash-attention, sharding, partition-spec
Unexpected mask shape: {mask.shape}
validation error jax, splash-attention, mask, shape-validation
Invalid shape for other: {other.shape}, expected: {self.shap
exception error jax, splash-attention, mask, shape-mismatch
chunk_size must be positive
validation error jax, splash-attention, mask, chunked-attention, argument-validation
Masks must have the same shape
validation error jax, splash-attention, mask, shape-mismatch
Unsupported empty tuple of masks
validation error jax, splash-attention, mask, empty-collection, validation
Unexpected mask shape, got: {mask.shape}, expected: {shape}
validation error jax, splash-attention, mask, shape-mismatch, multi-head
masks should be of type Mask
validation error jax, splash-attention, mask, type-validation
Nesting MultiHeadMasks is not supported
validation error jax, splash-attention, mask, multi-head, nesting
{q_block_size=} should divide {q_seq_len_per_shard=}.
validation error jax, tpu, splash-attention, shape-mismatch, pallas
{head_shards=} should divide {head_count=}.
validation error jax, tpu, splash-attention, sharding, shape-mismatch
Expected int32 input, but got {array.dtype}.
validation error jax, dtype, splash-attention, tpu, int64
Expected a 3-dim mask, instead got: {mask.shape=}
validation error jax, splash-attention, mask, rank-mismatch, tpu
bfloat16 support not implemented for LSTM
exception error jax, lstm, rnn, bfloat16, cudnn, gpu
Unexpected precision specifier value {precision}
exception error jax, lstm, rnn, invalid-argument, precision
`seq_lengths` can only be int32.
exception error jax, lstm, rnn, dtype, int32
Dropout not supported in LSTM reference because we cannot de
exception error jax, lstm, rnn, dropout, not-implemented
Unsupported shape: {shape}
exception error jax, sparse, coo, gpu, shape-validation
todense_transpose for {type(obj)}
exception error jax, sparse, autodiff, todense, csr
sparse_format={sparse_format!r} not recognized; must be one
exception error jax, sparse, argument-validation, format-string
batch_dims must be None or satisfy 0 < dim < n_batch. Got {b
exception error jax, sparse, bcoo, vmap, batching
data batch dimensions not compatible for {data.shape=}, {sha
exception error jax, sparse, bcoo, shape-validation, constructor
Invalid {data.shape=} for {nse=}, {n_batch=}, {n_dense=}
exception error jax, sparse, bcoo, shape-validation, nse
indices batch dimensions not compatible for {indices.shape=}
exception error jax, sparse, bcoo, indices, shape-validation
Invalid ={indices.shape=} for {nse=}, {n_batch=}, {n_dense=}
exception error jax, sparse, bcoo, indices, invariant
Expected 0 < bdim <= n_batch; got {bdim=}, {n_batch=}
exception error jax, sparse, bcoo, vmap, fromdense
First argument to bcoo_extract should be a BCOO array. Got {
exception error jax, sparse, type-check, bcoo
shape mismatch: {sparr.shape=} {a.shape=}
exception error jax, sparse, shape-validation, bcoo
transpose of bcoo_extract with assume_unique=False
exception error jax, sparse, bcoo, autodiff, transpose
bcoo_extract with unequal batch dimensions.
exception error jax, sparse, bcoo, vmap, batching
{batch_dims=} out of range for indices with {n_batch=}
exception error jax, sparse, bcoo, vmap, batching
transpose permutation must be a tuple/list/ndarray, got {typ
exception error jax, sparse, bcoo, transpose, type-check
transpose permutation isn't a permutation of operand dimensi
exception error jax, sparse, bcoo, transpose, permutation
transpose permutation cannot permute batch axes with non-bat
exception error jax, sparse, bcoo, transpose, permutation, batching
bcoo_slice: input should be BCOO array, got type(mat)={type(
exception error jax, sparse, bcoo, type-error, slice
bcoo_slice: indices must have size mat.ndim={mat.ndim}
exception error jax, sparse, bcoo, validation, shape-mismatch
len(strides) = {len(strides)}; expected {mat.ndim}
exception error jax, sparse, bcoo, strides, validation
strides must be a sequence of positive integers; got {stride
exception error jax, sparse, bcoo, strides, positive-value-required
bcoo_slice: invalid indices. Got {start_indices=}, {limit_in
exception error jax, sparse, bcoo, index-out-of-bounds, slice
bcoo_dynamic_slice: indices must have size mat.ndim={mat.ndi
exception error jax, sparse, bcoo, validation, shape-mismatch
slice_sizes must be less than or equal to operand shape, got
exception error jax, sparse, bcoo, window-too-large, dynamic-slice
bcoo_multiply_sparse: arrays must have same number of dimens
exception error jax, sparse, bcoo, shape-mismatch, elementwise-multiply
bcoo_multiply_sparse: arrays with differing numbers of dense
exception error jax, sparse, bcoo, not-implemented, dense-dimensions
BCSR sparse.empty: must have 2 sparse dimensions.
exception error jax, sparse, bcsr, shape-validation
BSCR.from_bcoo requires n_sparse=2; got {arr.n_sparse=}
exception error jax, sparse, bcsr, bcoo, format-conversion
BCSR from_scipy_sparse with nonzero n_dense/n_batch.
exception error jax, sparse, bcsr, scipy, not-implemented
BCSR from_scipy_sparse requires 2D array; {mat.ndim}D is giv
exception error jax, sparse, bcsr, scipy, input-validation
Cannot map in_axis={axis} for BCSR array with n_batch={val.n
exception error jax, sparse, bcsr, vmap, batching
BCSR: cannot add out_axis={axis} for BCSR array with n_batch
exception error jax, sparse, bcsr, vmap, out-axes
COO must have ndim=2; got {shape=}
exception error jax, sparse, coo, shape-validation
axes argument to transpose()
exception error jax, sparse, coo, transpose, not-implemented
COO.tree_unflatten: invalid {aux_data=}
exception error jax, sparse, coo, pytree, serialization
matmul between two sparse objects.
exception error jax, sparse, coo, matmul, not-implemented
matmul with object of shape {other.shape}
exception error jax, sparse, coo, matmul, shape-validation
CSR must have ndim=2; got {shape=}
exception error jax, sparse, csr, shape-validation
matmul between two sparse objects.
exception error jax, sparse, csr, matmul, not-implemented
matmul with object of shape {other.shape}
exception error jax, sparse, csr, matmul, shape-validation
CSR.tree_unflatten: invalid {aux_data=}
exception error jax, sparse, csr, pytree, serialization
CSC must have ndim=2; got {shape=}
exception error jax, sparse, csc, shape-validation
CSC.tree_unflatten: invalid {aux_data=}
exception error jax, sparse, pytree, csc, serialization
Cannot transpose with respect to sparse indices
exception error jax, sparse, csr, autodiff, gradient
must have search dim > 0, got {k}
exception error jax, sparse, lobpcg, eigensolver, validation
expected search dim * 5 < matrix dim (got {k * 5}, {n})
exception error jax, sparse, lobpcg, validation
A, X must have same dtypes (were {test_output.dtype}, {dt})
exception error jax, sparse, lobpcg, dtype
A must be ({n}, {n}) matrix A, got output {s}
exception error jax, sparse, lobpcg, shape-mismatch
data types do not match: {data.dtype=} {b.dtype=}
exception error jax, sparse, spsolve, dtype
index arrays must be integer typed; got {indices.dtype=} {in
validation error jax, sparse, spsolve, dtype, indices
Arrays must be one-dimensional. Got {data.shape=} {indices.s
validation error jax, sparse, spsolve, shape, dimensions
Invalid CSR buffer sizes: {data.shape=} {indices.shape=} {in
validation error jax, sparse, spsolve, csr, buffer-sizes
{reorder=} not valid, must be one of [1, 2, 3, 4]
validation error jax, sparse, spsolve, gpu, parameter-validation
Invalid {n_batch=}, {n_dense=} for {shape=}
validation error jax, sparse, bcoo, random, validation
got {nse=}, expected to be between 0 and {sparse_size}
validation error jax, sparse, bcoo, random, nse
{indices_dtype=} does not have enough range to generate spar
validation error jax, sparse, bcoo, random, index-overflow
sparse rule for {primitive} is not implemented because it wo
exception error jax, sparse, sparsify, lax, not-implemented
sparse rule for {primitive} is not implemented.
exception error jax, sparse, sparsify, not-implemented, lax
sparse rule for {lax.integer_pow_p} with non-positive expone
exception error jax, sparse, sparsify, power, densifying
Addition between sparse matrices of different shapes.
exception error jax, sparse, sparsify, addition, shape-mismatch
Addition between sparse matrices with different batch/dense
exception error jax, sparse, sparsify, addition, batch-dims
Error reading persistent compilation cache entry for '{cache
console warning jax, compilation-cache, persistence, io
Error reading persistent compilation cache entry for '{modul
console warning jax, compilation-cache, deserialization, io
Error writing persistent compilation cache entry for '{modul
console warning jax, compilation-cache, io, permissions
TRACER_LEAK_DEBUGGER_WARNING
console warning jax, tracer-leak, debugging, memory
JAX_SLICE_INDEX has been deprecated. Please use JAX_PARTITIO
console warning jax, distributed, environment-variable, deprecation
`slice_index` has been deprecated. Please use `partition_ind
console warning jax, distributed, api-rename, deprecation
Explicitly requested dtype {}{} is not available, and will b
console warning jax, dtype, precision, x64
When at least one mesh axis of `pred` is in auto mode, calli
console warning jax, distributed, spmd, sharding, error-check
cuDNN not found.
exception error jax, cuda, cudnn, gpu, shared-library
Argument to register_custom_call_partitioner was not a pjrt_
validation error jax, xla, custom-call, capsule, validation
This operation is implemented for a PjRt-compatible backend
exception error jax, dlpack, ifrt, pjrt, interop
DLPack is supported for PjRt-compatible backends only.
exception error jax, dlpack, ifrt, pjrt, interop
DLPack is only supported for devices addressable by the curr
exception error jax, dlpack, device, multi-process
dtype parameter is not supported by Buffer.__array__.
validation error jax, ffi, numpy, dtype
Buffer.__array__ with copy=True is not supported.
validation error jax, ffi, numpy, copy
Buffer.__dlpack__ with copy=True is not supported.
validation error jax, ffi, dlpack, copy
Unsupported index dtype: %s
validation error jax, cuda, cusparse, sparse, dtype
Unsupported data dtype: %s
validation error jax, gpu, sparse, dtype
Init is not implemented
exception error jax, collectives, cpu, api-misuse
Finalize is not implemented
exception error jax, collectives, cpu, api-misuse
make_gloo_tcp_collectives only implemented for linux and mac
exception error jax, gloo, collectives, platform-support
make_mpi_collectives is not implemented for Windows
exception error jax, mpi, collectives, platform-support
c_api argument to load_pjrt_plugin is not a pjrt_c_api capsu
validation error jax, pjrt, plugin, capsule, validation
Argument to get_c_api_topology was not a pjrt_c_api capsule.
validation error jax, pjrt, topology, capsule, validation
Argument to get_c_api_topology contained a null pointer.
validation error jax, pjrt, topology, capsule, null-pointer
get_topology_for_devices requires >= 1 devices.
validation error jax, topology, devices, validation
Expected tuple, got %s.
validation error jax, pytree, type-mismatch, tuple
Tuple arity mismatch: %d != %d; tuple: %s.
validation error jax, pytree, arity-mismatch, tuple
Expected list, got %s.
validation error jax, pytree, type-mismatch, list
List arity mismatch: %d != %d; list: %s.
validation error jax, pytree, arity-mismatch, list
Expected dict, got %s.
validation error jax, pytree, type-mismatch, dict
Dict key mismatch; expected keys: %s; present keys: %s.
validation error jax, pytree, dict, key-mismatch
Expected named tuple, got %s.
validation error jax, pytree, type-mismatch, namedtuple
Malformed pytree proto (invalid node type)
validation error jax, pytree, serialization, protobuf
Could not find type: %s.
exception error jax, pytree, pickle, registry
Pytree serialization too large to deserialize.
exception error jax, pytree, protobuf, size-limit
Could not deserialize PyTreeDefProto.
exception error jax, pytree, protobuf, corrupt-data
Malformed pickled PyTreeDef, expected 2-tuple
exception error jax, pytree, pickle
Malformed pickled SequenceKey, expected 1-tuple
exception error jax, pytree, pickle
Malformed pickled DictKey, expected 1-tuple
exception error jax, pytree, pickle
Malformed pickled GetAttrKey, expected 1-tuple
exception error jax, pytree, pickle
Malformed pickled FlattenedIndexKey, expected 1-tuple
exception error jax, pytree, pickle
Could not find memory addressable by device %s. Device %s ca
validation error jax, sharding, memory-kind, tpu
Got invalid memory kind: %s. Valid memory kinds are: %s
validation error jax, sharding, memory-kind
Invalid StrongLRUCache pickle version, got %d, expected 1
validation error jax, lru-cache, pickle, versioning
Cannot enter TracebackScope recursively.
validation error jax, traceback, context-manager
code argument must be a code object
validation error jax, traceback, type-check
Invalid pipeline stage
validation error jax, xla, compiler-passes, validation
c_api argument to register_xla_transform_c_api is not a pjrt
validation error jax, xla, pjrt, pycapsule
Cannot register XLA transform '%s': PJRT plugin does not sup
exception error jax, xla, pjrt, plugin-capabilities
c_api argument to clear_xla_transform_c_api is not a pjrt_c_
validation error jax, xla, pjrt, pycapsule