jax.lax module

jax.lax is a library of primitives operations that underpins libraries such as jax.numpy. Transformation rules, such as JVP and batching rules, are typically defined as transformations on jax.lax primitives.

Many of the primitives are thin wrappers around equivalent XLA operations, described by the XLA operation semantics documentation. In a few cases JAX diverges from XLA, usually to ensure that the set of operations is closed under the operation of JVP and transpose rules.

Where possible, prefer to use libraries such as jax.numpy instead of using jax.lax directly. The jax.numpy API follows NumPy, and is therefore more stable and less likely to change than the jax.lax API.

Operators

abs(x)	Elementwise absolute value: \(|x|\).
acos(x)	Elementwise arc cosine: \(\mathrm{acos}(x)\).
acosh(x)	Elementwise inverse hyperbolic cosine: \(\mathrm{acosh}(x)\).
add(x, y)	Elementwise addition: \(x + y\).
after_all(*operands)	Merges one or more XLA token values.
approx_max_k(operand, k[, ...])	Returns max k values and their indices of the operand in an approximate manner.
approx_min_k(operand, k[, ...])	Returns min k values and their indices of the operand in an approximate manner.
argmax(operand, axis, index_dtype)	Computes the index of the maximum element along axis.
argmin(operand, axis, index_dtype)	Computes the index of the minimum element along axis.
asin(x)	Elementwise arc sine: \(\mathrm{asin}(x)\).
asinh(x)	Elementwise inverse hyperbolic sine: \(\mathrm{asinh}(x)\).
atan(x)	Elementwise arc tangent: \(\mathrm{atan}(x)\).
atan2(x, y)	Elementwise two-term arc tangent: \(\mathrm{atan}({x \over y})\).
atanh(x)	Elementwise inverse hyperbolic tangent: \(\mathrm{atanh}(x)\).
batch_matmul(lhs, rhs[, precision])	Batch matrix multiplication.
bessel_i0e(x)	Exponentially scaled modified Bessel function of order 0: \(\mathrm{i0e}(x) = e^{-|x|} \mathrm{i0}(x)\)
bessel_i1e(x)	Exponentially scaled modified Bessel function of order 1: \(\mathrm{i1e}(x) = e^{-|x|} \mathrm{i1}(x)\)
betainc(a, b, x)	Elementwise regularized incomplete beta integral.
bitcast_convert_type(operand, new_dtype)	Elementwise bitcast.
bitwise_and(x, y)	Elementwise AND: \(x \wedge y\).
bitwise_not(x)	Elementwise NOT: \(\neg x\).
bitwise_or(x, y)	Elementwise OR: \(x \vee y\).
bitwise_xor(x, y)	Elementwise exclusive OR: \(x \oplus y\).
population_count(x)	Elementwise popcount, count the number of set bits in each element.
broadcast(operand, sizes[, out_sharding])	Broadcasts an array, adding new leading dimensions
broadcast_in_dim(operand, shape, ...[, ...])	Wraps XLA's BroadcastInDim operator.
broadcast_shapes(*shapes)	Returns the shape that results from NumPy broadcasting of shapes.
broadcast_to_rank(x, rank)	Adds leading dimensions of 1 to give x rank rank.
broadcasted_iota(dtype, shape, dimension[, ...])	Convenience wrapper around iota.
cbrt(x)	Elementwise cube root: \(\sqrt[3]{x}\).
ceil(x)	Elementwise ceiling: \(\left\lceil x \right\rceil\).
clamp(min, x, max)	Elementwise clamp.
clz(x)	Elementwise count-leading-zeros.
collapse(operand, start_dimension[, ...])	Collapses dimensions of an array into a single dimension.
complex(x, y)	Elementwise make complex number: \(x + jy\).
composite(decomposition, name[, version])	Composite with semantics defined by the decomposition function.
concatenate(operands, dimension)	Concatenates a sequence of arrays along dimension.
conj(x)	Elementwise complex conjugate function: \(\overline{x}\).
conv(lhs, rhs, window_strides, padding[, ...])	Convenience wrapper around conv_general_dilated.
convert_element_type(operand, new_dtype)	Elementwise cast.
conv_dimension_numbers(lhs_shape, rhs_shape, ...)	Converts convolution dimension_numbers to a ConvDimensionNumbers.
conv_general_dilated(lhs, rhs, ...[, ...])	General n-dimensional convolution operator, with optional dilation.
conv_general_dilated_local(lhs, rhs, ...[, ...])	General n-dimensional unshared convolution operator with optional dilation.
conv_general_dilated_patches(lhs, ...[, ...])	Extract patches subject to the receptive field of conv_general_dilated.
conv_transpose(lhs, rhs, strides, padding[, ...])	Convenience wrapper for calculating the N-d convolution "transpose".
conv_with_general_padding(lhs, rhs, ...[, ...])	Convenience wrapper around conv_general_dilated.
cos(x)	Elementwise cosine: \(\mathrm{cos}(x)\).
cosh(x)	Elementwise hyperbolic cosine: \(\mathrm{cosh}(x)\).
cumlogsumexp(operand[, axis, reverse])	Computes a cumulative logsumexp along axis.
cummax(operand[, axis, reverse])	Computes a cumulative maximum along axis.
cummin(operand[, axis, reverse])	Computes a cumulative minimum along axis.
cumprod(operand[, axis, reverse])	Computes a cumulative product along axis.
cumsum(operand[, axis, reverse])	Computes a cumulative sum along axis.
digamma(x)	Elementwise digamma: \(\psi(x)\).
div(x, y)	Elementwise division: \(x \over y\).
dot(lhs, rhs[, precision, ...])	Vector/vector, matrix/vector, and matrix/matrix multiplication.
dot_general(lhs, rhs, dimension_numbers[, ...])	General dot product/contraction operator.
dynamic_index_in_dim(operand, index[, axis, ...])	Convenience wrapper around dynamic_slice to perform int indexing.
dynamic_slice(operand, start_indices, ...[, ...])	Wraps XLA's DynamicSlice operator.
dynamic_slice_in_dim(operand, start_index, ...)	Convenience wrapper around lax.dynamic_slice() applied to one dimension.
dynamic_update_index_in_dim(operand, update, ...)	Convenience wrapper around dynamic_update_slice() to update a slice of size 1 in a single axis.
dynamic_update_slice(operand, update, ...[, ...])	Wraps XLA's DynamicUpdateSlice operator.
dynamic_update_slice_in_dim(operand, update, ...)	Convenience wrapper around dynamic_update_slice() to update a slice in a single axis.
eq(x, y)	Elementwise equals: \(x = y\).
erf(x)	Elementwise error function: \(\mathrm{erf}(x)\).
erfc(x)	Elementwise complementary error function: \(\mathrm{erfc}(x) = 1 - \mathrm{erf}(x)\).
erf_inv(x)	Elementwise inverse error function: \(\mathrm{erf}^{-1}(x)\).
exp(x)	Elementwise exponential: \(e^x\).
exp2(x)	Elementwise base-2 exponential: \(2^x\).
expand_dims(array, dimensions)	Insert any number of size 1 dimensions into an array.
expm1(x)	Elementwise \(e^{x} - 1\).
fft(x, fft_type, fft_lengths)
floor(x)	Elementwise floor: \(\left\lfloor x \right\rfloor\).
full(shape, fill_value[, dtype, sharding])	Returns an array of shape filled with fill_value.
full_like(x, fill_value[, dtype, shape, ...])	Create a full array like np.full based on the example array x.
gather(operand, start_indices, ...[, ...])	Gather operator.
ge(x, y)	Elementwise greater-than-or-equals: \(x \geq y\).
gt(x, y)	Elementwise greater-than: \(x > y\).
igamma(a, x)	Elementwise regularized incomplete gamma function.
igammac(a, x)	Elementwise complementary regularized incomplete gamma function.
imag(x)	Elementwise extract imaginary part: \(\mathrm{Im}(x)\).
index_in_dim(operand, index[, axis, keepdims])	Convenience wrapper around lax.slice() to perform int indexing.
index_take(src, idxs, axes)
integer_pow(x, y)	Elementwise power: \(x^y\), where \(y\) is a static integer.
iota(dtype, size)	Wraps XLA's Iota operator.
is_finite(x)	Elementwise \(\mathrm{isfinite}\).
le(x, y)	Elementwise less-than-or-equals: \(x \leq y\).
lgamma(x)	Elementwise log gamma: \(\mathrm{log}(\Gamma(x))\).
log(x)	Elementwise natural logarithm: \(\mathrm{log}(x)\).
log1p(x)	Elementwise \(\mathrm{log}(1 + x)\).
logistic(x)	Elementwise logistic (sigmoid) function: \(\frac{1}{1 + e^{-x}}\).
lt(x, y)	Elementwise less-than: \(x < y\).
max(x, y)	Elementwise maximum: \(\mathrm{max}(x, y)\)
min(x, y)	Elementwise minimum: \(\mathrm{min}(x, y)\)
mul(x, y)	Elementwise multiplication: \(x \times y\).
ne(x, y)	Elementwise not-equals: \(x \neq y\).
neg(x)	Elementwise negation: \(-x\).
nextafter(x1, x2)	Returns the next representable value after x1 in the direction of x2.
optimization_barrier(operand, /)	Prevents the compiler from moving operations across the barrier.
pad(operand, padding_value, padding_config)	Applies low, high, and/or interior padding to an array.
platform_dependent(*args[, default])	Stages out platform-specific code.
polygamma(m, x)	Elementwise polygamma: \(\psi^{(m)}(x)\).
population_count(x)	Elementwise popcount, count the number of set bits in each element.
pow(x, y)	Elementwise power: \(x^y\).
random_gamma_grad(a, x)	Elementwise derivative of samples from Gamma(a, 1).
real(x)	Elementwise extract real part: \(\mathrm{Re}(x)\).
reciprocal(x)	Elementwise reciprocal: \(1 \over x\).
reduce(operands, init_values, computation, ...)	Wraps XLA's Reduce operator.
reduce_and(operand, axes)	Compute the bitwise AND of elements over one or more array axes.
reduce_max(operand, axes)	Compute the maximum of elements over one or more array axes.
reduce_min(operand, axes)	Compute the minimum of elements over one or more array axes.
reduce_or(operand, axes)	Compute the bitwise OR of elements over one or more array axes.
reduce_precision(operand, exponent_bits, ...)	Wraps XLA's ReducePrecision operator.
reduce_prod(operand, axes)	Compute the product of elements over one or more array axes.
reduce_sum(operand, axes)	Compute the sum of elements over one or more array axes.
reduce_window(operand, init_value, ...[, ...])	Wraps XLA's ReduceWindowWithGeneralPadding operator.
reduce_xor(operand, axes)	Compute the bitwise XOR of elements over one or more array axes.
rem(x, y)	Elementwise remainder: \(x \bmod y\).
reshape(operand, new_sizes[, dimensions, ...])	Wraps XLA's Reshape operator.
rev(operand, dimensions)	Wraps XLA's Rev operator.
rng_bit_generator(key, shape[, dtype, algorithm])	Stateless PRNG bit generator.
rng_uniform(a, b, shape)	Stateful PRNG generator.
round(x[, rounding_method])	Elementwise round.
rsqrt(x)	Elementwise reciprocal square root: \(1 \over \sqrt{x}\).
scatter(operand, scatter_indices, updates, ...)	Scatter-update operator.
scatter_add(operand, scatter_indices, ...[, ...])	Scatter-add operator.
scatter_apply(operand, scatter_indices, ...)	Scatter-apply operator.
scatter_max(operand, scatter_indices, ...[, ...])	Scatter-max operator.
scatter_min(operand, scatter_indices, ...[, ...])	Scatter-min operator.
scatter_mul(operand, scatter_indices, ...[, ...])	Scatter-multiply operator.
shift_left(x, y)	Elementwise left shift: \(x \ll y\).
shift_right_arithmetic(x, y)	Elementwise arithmetic right shift: \(x \gg y\).
shift_right_logical(x, y)	Elementwise logical right shift: \(x \gg y\).
sign(x)	Elementwise sign.
sin(x)	Elementwise sine: \(\mathrm{sin}(x)\).
sinh(x)	Elementwise hyperbolic sine: \(\mathrm{sinh}(x)\).
slice(operand, start_indices, limit_indices)	Wraps XLA's Slice operator.
slice_in_dim(operand, start_index, limit_index)	Convenience wrapper around lax.slice() applying to only one dimension.
sort()	Wraps XLA's Sort operator.
sort_key_val(keys, values[, dimension, ...])	Sorts keys along dimension and applies the same permutation to values.
split(operand, sizes[, axis])	Splits an array along axis.
sqrt(x)	Elementwise square root: \(\sqrt{x}\).
square(x)	Elementwise square: \(x^2\).
squeeze(array, dimensions)	Squeeze any number of size 1 dimensions from an array.
sub(x, y)	Elementwise subtraction: \(x - y\).
tan(x)	Elementwise tangent: \(\mathrm{tan}(x)\).
tanh(x)	Elementwise hyperbolic tangent: \(\mathrm{tanh}(x)\).
top_k(operand, k)	Returns top k values and their indices along the last axis of operand.
transpose(operand, permutation)	Wraps XLA's Transpose operator.
zeros_like_array(x)
zeta(x, q)	Elementwise Hurwitz zeta function: \(\zeta(x, q)\)

Control flow operators

associative_scan(fn, elems[, reverse, axis])	Performs a scan with an associative binary operation, in parallel.
cond(pred, true_fun, false_fun, *operands[, ...])	Conditionally apply true_fun or false_fun.
fori_loop(lower, upper, body_fun, init_val, *)	Loop from lower to upper by reduction to jax.lax.while_loop().
map(f, xs, *[, batch_size])	Map a function over leading array axes.
scan(f, init[, xs, length, reverse, unroll, ...])	Scan a function over leading array axes while carrying along state.
select(pred, on_true, on_false)	Selects between two branches based on a boolean predicate.
select_n(which, *cases)	Selects array values from multiple cases.
switch(index, branches, *operands[, operand])	Apply exactly one of the branches given by index.
while_loop(cond_fun, body_fun, init_val)	Call body_fun repeatedly in a loop while cond_fun is True.

Custom gradient operators

stop_gradient(x)	Stops gradient computation.
custom_linear_solve(matvec, b, solve[, ...])	Perform a matrix-free linear solve with implicitly defined gradients.
custom_root(f, initial_guess, solve, ...[, ...])	Differentiably solve for the roots of a function.

Parallel operators

all_gather(x, axis_name, *[, ...])	Gather values of x across all replicas.
all_to_all(x, axis_name, split_axis, ...[, ...])	Materialize the mapped axis and map a different axis.
psum(x, axis_name, *[, axis_index_groups])	Compute an all-reduce sum on x over the pmapped axis axis_name.
psum_scatter(x, axis_name, *[, ...])	Like psum(x, axis_name) but each device retains only part of the result.
pmax(x, axis_name, *[, axis_index_groups])	Compute an all-reduce max on x over the pmapped axis axis_name.
pmin(x, axis_name, *[, axis_index_groups])	Compute an all-reduce min on x over the pmapped axis axis_name.
pmean(x, axis_name, *[, axis_index_groups])	Compute an all-reduce mean on x over the pmapped axis axis_name.
ppermute(x, axis_name, perm)	Perform a collective permutation according to the permutation perm.
pshuffle(x, axis_name, perm)	Convenience wrapper of jax.lax.ppermute with alternate permutation encoding
pswapaxes(x, axis_name, axis, *[, ...])	Swap the pmapped axis axis_name with the unmapped axis axis.
axis_index(axis_name)	Return the index along the mapped axis axis_name.

Sharding-related operators

with_sharding_constraint(x, shardings)

Mechanism to constrain the sharding of an Array inside a jitted computation

Linear algebra operators (jax.lax.linalg)

cholesky(x, *[, symmetrize_input])	Cholesky decomposition.
cholesky_update(r_matrix, w_vector)	Cholesky rank-1 update.
eig(x, *[, compute_left_eigenvectors, ...])	Eigendecomposition of a general matrix.
eigh(x, *[, lower, symmetrize_input, ...])	Eigendecomposition of a Hermitian matrix.
hessenberg(a)	Reduces a square matrix to upper Hessenberg form.
householder_product(a, taus)	Product of elementary Householder reflectors.
lu(x)	LU decomposition with partial pivoting.
lu_pivots_to_permutation(pivots, ...)	Converts the pivots (row swaps) returned by LU to a permutation.
qdwh(x, *[, is_hermitian, max_iterations, ...])	QR-based dynamically weighted Halley iteration for polar decomposition.
qr()	QR decomposition.
schur(x, *[, compute_schur_vectors, ...])	Schur decomposition.
svd()	Singular value decomposition.
SvdAlgorithm(value)	Enum for SVD algorithm.
symmetric_product(a_matrix, c_matrix, *[, ...])	Symmetric product.
triangular_solve(a, b, *[, left_side, ...])	Triangular solve.
tridiagonal(a, *[, lower])	Reduces a symmetric/Hermitian matrix to tridiagonal form.
tridiagonal_solve(dl, d, du, b)	Computes the solution of a tridiagonal linear system.

Argument classes

class jax.lax.ConvDimensionNumbers(lhs_spec, rhs_spec, out_spec)

Describes batch, spatial, and feature dimensions of a convolution.

Parameters:

• lhs_spec (Sequence[int]) – a tuple of nonnegative integer dimension numbers containing (batch dimension, feature dimension, spatial dimensions…).

• rhs_spec (Sequence[int]) – a tuple of nonnegative integer dimension numbers containing (out feature dimension, in feature dimension, spatial dimensions…).

• out_spec (Sequence[int]) – a tuple of nonnegative integer dimension numbers containing (batch dimension, feature dimension, spatial dimensions…).

jax.lax.ConvGeneralDilatedDimensionNumbers

alias of tuple[str, str, str] | ConvDimensionNumbers | None

class jax.lax.DotAlgorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count=1, rhs_component_count=1, num_primitive_operations=1, allow_imprecise_accumulation=False)

Specify the algorithm used for computing dot products.

When used to specify the precision input to dot(), dot_general(), and other dot product functions, this data structure is used for controlling the properties of the algorithm used for computing the dot product. This API controls the precision used for the computation, and allows users to access hardware-specific accelerations.

Support for these algorithms is platform dependent, and using an unsupported algorithm will raise a Python exception when the computation is compiled. The algorithms that are known to be supported on at least some platforms are listed in the DotAlgorithmPreset enum, and these are a good starting point for experimenting with this API.

A “dot algorithm” is specified by the following parameters:

• lhs_precision_type and rhs_precision_type, the data types that the LHS and RHS of the operation are rounded to.

• accumulation_type the data type used for accumulation.

• lhs_component_count, rhs_component_count, and num_primitive_operations apply to algorithms that decompose the LHS and/or RHS into multiple components and execute multiple operations on those values, usually to emulate a higher precision. For algorithms with no decomposition, these values should be set to 1.

• allow_imprecise_accumulation to specify if accumulation in lower precision is permitted for some steps (e.g. CUBLASLT_MATMUL_DESC_FAST_ACCUM).

The StableHLO spec for the dot operation doesn’t require that the precision types be the same as the storage types for the inputs or outputs, but some plaforms may require that these types match. Furthermore, the return type of dot_general() is always defined by the accumulation_type parameter of the input algorithm, if specified.

Examples

Accumulate two 16-bit floats using a 32-bit float accumulator:

>>> algorithm = DotAlgorithm(
...     lhs_precision_type=np.float16,
...     rhs_precision_type=np.float16,
...     accumulation_type=np.float32,
... )
>>> lhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)
>>> rhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)
>>> dot(lhs, rhs, precision=algorithm)  
array([ 1.,  4.,  9., 16.], dtype=float16)

Or, equivalently, using a preset:

>>> algorithm = DotAlgorithmPreset.F16_F16_F32
>>> dot(lhs, rhs, precision=algorithm)  
array([ 1.,  4.,  9., 16.], dtype=float16)

Presets can also be specified by name:

>>> dot(lhs, rhs, precision="F16_F16_F32")  
array([ 1.,  4.,  9., 16.], dtype=float16)

The preferred_element_type parameter can be used to return the output without downcasting the accumulation type:

>>> dot(lhs, rhs, precision="F16_F16_F32", preferred_element_type=np.float32)  
array([ 1.,  4.,  9., 16.], dtype=float32)

Parameters:

• lhs_precision_type (DTypeLike)

• rhs_precision_type (DTypeLike)

• accumulation_type (DTypeLike)

• lhs_component_count (int)

• rhs_component_count (int)

• num_primitive_operations (int)

• allow_imprecise_accumulation (bool)

class jax.lax.DotAlgorithmPreset(value)

An enum of known algorithms for computing dot products.

This Enum provides a named set of DotAlgorithm objects that are known to be supported on at least platform. See the DotAlgorithm documentation for more details about the behavior of these algorithms.

An algorithm can be selected from this list when calling dot(), dot_general(), or most other JAX dot product functions, by passing either a member of this Enum or it’s name as a string using the precision argument.

For example, users can specify the preset using this Enum directly:

>>> lhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)
>>> rhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)
>>> algorithm = DotAlgorithmPreset.F16_F16_F32
>>> dot(lhs, rhs, precision=algorithm)  
array([ 1.,  4.,  9., 16.], dtype=float16)

or, equivalently, they can be specified by name:

>>> dot(lhs, rhs, precision="F16_F16_F32")  
array([ 1.,  4.,  9., 16.], dtype=float16)

The names of the presets are typically LHS_RHS_ACCUM where LHS and RHS are the element types of the lhs and rhs inputs respectively, and ACCUM is the element type of the accumulator. Some presets have an extra suffix, and the meaning of each of these is documented below. The supported presets are:

DEFAULT = 1

An algorithm will be selected based on input and output types.

ANY_F8_ANY_F8_F32 = 2

Accepts any float8 input types and accumulates into float32.

ANY_F8_ANY_F8_F32_FAST_ACCUM = 3

Like ANY_F8_ANY_F8_F32, but using faster accumulation with the cost of lower accuracy.

ANY_F8_ANY_F8_ANY = 4

Like ANY_F8_ANY_F8_F32, but the accumulation type is controlled by preferred_element_type.

ANY_F8_ANY_F8_ANY_FAST_ACCUM = 5

Like ANY_F8_ANY_F8_F32_FAST_ACCUM, but the accumulation type is controlled by preferred_element_type.

F16_F16_F16 = 6

F16_F16_F32 = 7

BF16_BF16_BF16 = 8

BF16_BF16_F32 = 9

BF16_BF16_F32_X3 = 10

The _X3 suffix indicates that the algorithm uses 3 operations to emulate higher precision.

BF16_BF16_F32_X6 = 11

Like BF16_BF16_F32_X3, but using 6 operations instead of 3.

BF16_BF16_F32_X9 = 12

Like BF16_BF16_F32_X3, but using 9 operations instead of 3.

TF32_TF32_F32 = 13

TF32_TF32_F32_X3 = 14

The _X3 suffix indicates that the algorithm uses 3 operations to emulate higher precision.

F32_F32_F32 = 15

F64_F64_F64 = 16

property supported_lhs_types: tuple[DTypeLike, ...] | None

property supported_rhs_types: tuple[DTypeLike, ...] | None

property accumulation_type: DTypeLike | None

supported_output_types(lhs_dtype, rhs_dtype)

Parameters:

• lhs_dtype (DTypeLike)

• rhs_dtype (DTypeLike)

Return type:

tuple[DTypeLike, …] | None

class jax.lax.FftType(value)

Describes which FFT operation to perform.

FFT = 0

Forward complex-to-complex FFT.

IFFT = 1

Inverse complex-to-complex FFT.

IRFFT = 3

Inverse real-to-complex FFT.

RFFT = 2

Forward real-to-complex FFT.

class jax.lax.GatherDimensionNumbers(offset_dims, collapsed_slice_dims, start_index_map, operand_batching_dims=(), start_indices_batching_dims=())

Describes the dimension number arguments to an XLA’s Gather operator. See the XLA documentation for more details of what the dimension numbers mean.

Parameters:

• offset_dims (tuple[int, ...]) – the set of dimensions in the gather output that offset into an array sliced from operand. Must be a tuple of integers in ascending order, each representing a dimension number of the output.

• collapsed_slice_dims (tuple[int, ...]) – the set of dimensions i in operand that have slice_sizes[i] == 1 and that should not have a corresponding dimension in the output of the gather. Must be a tuple of integers in ascending order.

• start_index_map (tuple[int, ...]) – for each dimension in start_indices, gives the corresponding dimension in the operand that is to be sliced. Must be a tuple of integers with size equal to start_indices.shape[-1].

• operand_batching_dims (tuple[int, ...]) – the set of batching dimensions i in operand that have slice_sizes[i] == 1 and that should have a corresponding dimension in both the start_indices (at the same index in start_indices_batching_dims) and output of the gather. Must be a tuple of integers in ascending order.

• start_indices_batching_dims (tuple[int, ...]) – the set of batching dimensions i in start_indices that should have a corresponding dimension in both the operand (at the same index in operand_batching_dims) and output of the gather. Must be a tuple of integers (order is fixed based on correspondence with operand_batching_dims).

Unlike XLA’s GatherDimensionNumbers structure, index_vector_dim is implicit; there is always an index vector dimension and it must always be the last dimension. To gather scalar indices, add a trailing dimension of size 1.

class jax.lax.GatherScatterMode(value)

Describes how to handle out-of-bounds indices in a gather or scatter.

Possible values are:

CLIP:

Indices will be clamped to the nearest in-range value, i.e., such that the entire window to be gathered is in-range.

FILL_OR_DROP:

If any part of a gathered window is out of bounds, the entire window that is returned, even those elements that were otherwise in-bounds, will be filled with a constant. If any part of a scattered window is out of bounds, the entire window will be discarded.

PROMISE_IN_BOUNDS:

The user promises that indices are in bounds. No additional checking will be performed. In practice, with the current XLA implementation this means that out-of-bounds gathers will be clamped but out-of-bounds scatters will be discarded. Gradients will not be correct if indices are out-of-bounds.

class jax.lax.Precision(value)

Precision enum for lax matrix multiply related functions.

The device-dependent precision argument to JAX functions generally controls the tradeoff between speed and accuracy for array computations on accelerator backends, (i.e. TPU and GPU). Has no impact on CPU backends. This only has an effect on float32 computations, and does not affect the input/output datatypes. Members are:

DEFAULT:

Fastest mode, but least accurate. On TPU: performs float32 computations in bfloat16. On GPU: uses tensorfloat32 if available (e.g. on A100 and H100 GPUs), otherwise standard float32 (e.g. on V100 GPUs). Aliases: 'default', 'fastest'.

HIGH:

Slower but more accurate. On TPU: performs float32 computations in 3 bfloat16 passes. On GPU: uses tensorfloat32 where available, otherwise float32. Aliases: 'high'..

HIGHEST:

Slowest but most accurate. On TPU: performs float32 computations in 6 bfloat16. Aliases: 'highest'. On GPU: uses float32.

jax.lax.PrecisionLike

alias of None | str | Precision | tuple[str, str] | tuple[Precision, Precision] | DotAlgorithm | DotAlgorithmPreset

class jax.lax.RandomAlgorithm(value)

Describes which PRNG algorithm to use for rng_bit_generator.

RNG_DEFAULT = 0

The platform’s default algorithm.

RNG_THREE_FRY = 1

The Threefry-2x32 PRNG algorithm.

RNG_PHILOX = 2

The Philox-4x32 PRNG algorithm.

class jax.lax.RoundingMethod(value)

Rounding strategies for handling halfway values (e.g., 0.5) in jax.lax.round().

AWAY_FROM_ZERO = 0

Rounds halfway values away from zero (e.g., 0.5 -> 1, -0.5 -> -1).

TO_NEAREST_EVEN = 1

Rounds halfway values to the nearest even integer. This is also known as “banker’s rounding” (e.g., 0.5 -> 0, 1.5 -> 2).

class jax.lax.ScatterDimensionNumbers(update_window_dims, inserted_window_dims, scatter_dims_to_operand_dims, operand_batching_dims=(), scatter_indices_batching_dims=())

Describes the dimension number arguments to an XLA’s Scatter operator. See the XLA documentation for more details of what the dimension numbers mean.

Parameters:

• update_window_dims (Sequence[int]) – the set of dimensions in the updates that are window dimensions. Must be a tuple of integers in ascending order, each representing a dimension number.

• inserted_window_dims (Sequence[int]) – the set of size 1 window dimensions that must be inserted into the shape of updates. Must be a tuple of integers in ascending order, each representing a dimension number of the output. These are the mirror image of collapsed_slice_dims in the case of gather.

• scatter_dims_to_operand_dims (Sequence[int]) – for each dimension in scatter_indices, gives the corresponding dimension in operand. Must be a sequence of integers with size equal to scatter_indices.shape[-1].

• operand_batching_dims (Sequence[int]) – the set of batching dimensions i in operand that should have a corresponding dimension in both the scatter_indices (at the same index in scatter_indices_batching_dims) and updates. Must be a tuple of integers in ascending order. These are the mirror image of operand_batching_dims in the case of gather.

• scatter_indices_batching_dims (Sequence[int]) – the set of batching dimensions i in scatter_indices that should have a corresponding dimension in both the operand (at the same index in operand_batching_dims) and output of the gather. Must be a tuple of integers (order is fixed based on correspondence with input_batching_dims). These are the mirror image of start_indices_batching_dims in the case of gather.

Unlike XLA’s ScatterDimensionNumbers structure, index_vector_dim is implicit; there is always an index vector dimension and it must always be the last dimension. To scatter scalar indices, add a trailing dimension of size 1.