jax.lax.shift_right_arithmetic

jax.lax.shift_right_arithmetic(x, y)

Elementwise arithmetic right shift: \(x \gg y\).

This function lowers directly to the stablehlo.shift_right_arithmetic operation.

Parameters:

• x (ArrayLike) – Input arrays. Must have matching integer dtypes. If neither is a scalar, x and y must have the same number of dimensions and be broadcast compatible.

• y (ArrayLike) – Input arrays. Must have matching integer dtypes. If neither is a scalar, x and y must have the same number of dimensions and be broadcast compatible.

Returns:

An array of the same dtype as x and y containing the element-wise arithmetic right shift of each pair of broadcasted entries.

Return type:

Array

See also

• jax.numpy.right_shift(): NumPy wrapper for this API when applied to signed integers, also accessible via the x >> y operator on JAX arrays with signed integer dtype.

• jax.lax.shift_left(): Elementwise left shift.

• jax.lax.shift_right_logical(): Elementwise logical right shift.