Operations#

abs(a, /, *[, stream])

Element-wise absolute value.

add(a, b[, stream])

Element-wise addition.

all(a, /[, axis, keepdims, stream])

An and reduction over the given axes.

allclose(a, b, /[, rtol, atol, equal_nan, ...])

Approximate comparison of two arrays.

any(a, /[, axis, keepdims, stream])

An or reduction over the given axes.

arange(-> array)

Generates ranges of numbers.

arccos(a, /, *[, stream])

Element-wise inverse cosine.

arccosh(a, /, *[, stream])

Element-wise inverse hyperbolic cosine.

arcsin(a, /, *[, stream])

Element-wise inverse sine.

arcsinh(a, /, *[, stream])

Element-wise inverse hyperbolic sine.

arctan(a, /, *[, stream])

Element-wise inverse tangent.

arctan2(a, b, /, *[, stream])

Element-wise inverse tangent of the ratio of two arrays.

arctanh(a, /, *[, stream])

Element-wise inverse hyperbolic tangent.

argmax(a, /[, axis, keepdims, stream])

Indices of the maximum values along the axis.

argmin(a, /[, axis, keepdims, stream])

Indices of the minimum values along the axis.

argpartition(a, /, kth[, axis, stream])

Returns the indices that partition the array.

argsort(a, /[, axis, stream])

Returns the indices that sort the array.

array_equal(a, b[, equal_nan, stream])

Array equality check.

atleast_1d(*arys[, stream])

Convert all arrays to have at least one dimension.

atleast_2d(*arys[, stream])

Convert all arrays to have at least two dimensions.

atleast_3d(*arys[, stream])

Convert all arrays to have at least three dimensions.

bitwise_and(a, b[, stream])

Element-wise bitwise and.

bitwise_or(a, b[, stream])

Element-wise bitwise or.

bitwise_xor(a, b[, stream])

Element-wise bitwise xor.

block_masked_mm(a, b, /, block_size, ...)

Matrix multiplication with block masking.

block_sparse_mm(a, b, /, lhs_indices, ...[, ...])

Matrix multiplication with matrix-level gather.

broadcast_to(a, /, shape, *[, stream])

Broadcast an array to the given shape.

ceil(a, /, *[, stream])

Element-wise ceil.

clip(a, /, a_min, a_max, *[, stream])

Clip the values of the array between the given minimum and maximum.

concatenate(arrays[, axis, stream])

Concatenate the arrays along the given axis.

conj(a, *[, stream])

Return the elementwise complex conjugate of the input.

conjugate(a, *[, stream])

Return the elementwise complex conjugate of the input.

convolve(a, v, /[, mode, stream])

The discrete convolution of 1D arrays.

conv1d(input, weight, /[, stride, padding, ...])

1D convolution over an input with several channels

conv2d(input, weight, /[, stride, padding, ...])

2D convolution over an input with several channels

conv_general(input, weight, /[, stride, ...])

General convolution over an input with several channels

cos(a, /, *[, stream])

Element-wise cosine.

cosh(a, /, *[, stream])

Element-wise hyperbolic cosine.

cummax(a, /[, axis, reverse, inclusive, stream])

Return the cumulative maximum of the elements along the given axis.

cummin(a, /[, axis, reverse, inclusive, stream])

Return the cumulative minimum of the elements along the given axis.

cumprod(a, /[, axis, reverse, inclusive, stream])

Return the cumulative product of the elements along the given axis.

cumsum(a, /[, axis, reverse, inclusive, stream])

Return the cumulative sum of the elements along the given axis.

degrees(a, /, *[, stream])

Convert angles from radians to degrees.

dequantize(w, /, scales, biases[, ...])

Dequantize the matrix w using the provided scales and biases and the group_size and bits configuration.

diag(a, /[, k, stream])

Extract a diagonal or construct a diagonal matrix.

diagonal(a[, offset, axis1, axis2, stream])

Return specified diagonals.

divide(a, b[, stream])

Element-wise division.

divmod(a, b[, stream])

Element-wise quotient and remainder.

equal(a, b[, stream])

Element-wise equality.

erf(a, /, *[, stream])

Element-wise error function.

erfinv(a, /, *[, stream])

Element-wise inverse of erf().

exp(a, /, *[, stream])

Element-wise exponential.

expm1(a, /, *[, stream])

Element-wise exponential minus 1.

expand_dims(a, /, axis, *[, stream])

Add a size one dimension at the given axis.

eye(n[, m, k, dtype, stream])

Create an identity matrix or a general diagonal matrix.

flatten(a, /[, start_axis, end_axis, stream])

Flatten an array.

floor(a, /, *[, stream])

Element-wise floor.

floor_divide(a, b[, stream])

Element-wise integer division.

full(shape, vals[, dtype, stream])

Construct an array with the given value.

greater(a, b[, stream])

Element-wise greater than.

greater_equal(a, b[, stream])

Element-wise greater or equal.

identity(n[, dtype, stream])

Create a square identity matrix.

inner(a, b, /, *[, stream])

Ordinary inner product of vectors for 1-D arrays, in higher dimensions a sum product over the last axes.

isclose(a, b, /[, rtol, atol, equal_nan, stream])

Returns a boolean array where two arrays are element-wise equal within a tolerance.

isinf(a[, stream])

Return a boolean array indicating which elements are +/- inifnity.

isnan(a[, stream])

Return a boolean array indicating which elements are NaN.

isneginf(a[, stream])

Return a boolean array indicating which elements are negative infinity.

isposinf(a[, stream])

Return a boolean array indicating which elements are positive infinity.

left_shift(a, b[, stream])

Element-wise left shift.

less(a, b[, stream])

Element-wise less than.

less_equal(a, b[, stream])

Element-wise less than or equal.

linspace(start, stop[, num, dtype, stream])

Generate num evenly spaced numbers over interval [start, stop].

load(file, /[, format, return_metadata, stream])

Load array(s) from a binary file.

log(a, /, *[, stream])

Element-wise natural logarithm.

log2(a, /, *[, stream])

Element-wise base-2 logarithm.

log10(a, /, *[, stream])

Element-wise base-10 logarithm.

log1p(a, /, *[, stream])

Element-wise natural log of one plus the array.

logaddexp(a, b, /, *[, stream])

Element-wise log-add-exp.

logical_not(a, /, *[, stream])

Element-wise logical not.

logical_and(a, b, /, *[, stream])

Element-wise logical and.

logical_or(a, b, /, *[, stream])

Element-wise logical or.

logsumexp(a, /[, axis, keepdims, stream])

A log-sum-exp reduction over the given axes.

matmul(a, b, /, *[, stream])

Matrix multiplication.

max(a, /[, axis, keepdims, stream])

A max reduction over the given axes.

maximum(a, b, /, *[, stream])

Element-wise maximum.

mean(a, /[, axis, keepdims, stream])

Compute the mean(s) over the given axes.

meshgrid(*arrays[, sparse, indexing, stream])

Generate multidimensional coordinate grids from 1-D coordinate arrays

min(a, /[, axis, keepdims, stream])

A min reduction over the given axes.

minimum(a, b, /, *[, stream])

Element-wise minimum.

moveaxis(a, /, source, destination, *[, stream])

Move an axis to a new position.

multiply(a, b[, stream])

Element-wise multiplication.

negative(a, /, *[, stream])

Element-wise negation.

not_equal(a, b[, stream])

Element-wise not equal.

ones(shape[, dtype, stream])

Construct an array of ones.

ones_like(a, /, *[, stream])

An array of ones like the input.

outer(a, b, /, *[, stream])

Compute the outer product of two 1-D arrays, if the array's passed are not 1-D a flatten op will be run beforehand.

partition(a, /, kth[, axis, stream])

Returns a partitioned copy of the array such that the smaller kth elements are first.

pad(a, pad_with[, constant_values, stream])

Pad an array with a constant value

prod(a, /[, axis, keepdims, stream])

An product reduction over the given axes.

quantize(w, /[, group_size, bits, stream])

Quantize the matrix w using bits bits per element.

quantized_matmul(x, w, /, scales, biases[, ...])

Perform the matrix multiplication with the quantized matrix w.

radians(a, /, *[, stream])

Convert angles from degrees to radians.

reciprocal(a, /, *[, stream])

Element-wise reciprocal.

repeat(array, repeats[, axis, stream])

Repeat an array along a specified axis.

reshape(a, /, shape, *[, stream])

Reshape an array while preserving the size.

right_shift(a, b[, stream])

Element-wise right shift.

round(a, /[, decimals, stream])

Round to the given number of decimals.

rsqrt(a, /, *[, stream])

Element-wise reciprocal and square root.

save(file, arr)

Save the array to a binary file in .npy format.

savez(file, *args, **kwargs)

Save several arrays to a binary file in uncompressed .npz format.

savez_compressed(file, *args, **kwargs)

Save several arrays to a binary file in compressed .npz format.

save_gguf(file, arrays, metadata)

Save array(s) to a binary file in .gguf format.

save_safetensors(file, arrays[, metadata])

Save array(s) to a binary file in .safetensors format.

sigmoid(a, /, *[, stream])

Element-wise logistic sigmoid.

sign(a, /, *[, stream])

Element-wise sign.

sin(a, /, *[, stream])

Element-wise sine.

sinh(a, /, *[, stream])

Element-wise hyperbolic sine.

softmax(a, /[, axis, stream])

Perform the softmax along the given axis.

sort(a, /[, axis, stream])

Returns a sorted copy of the array.

split(a, /, indices_or_sections[, axis, stream])

Split an array along a given axis.

sqrt(a, /, *[, stream])

Element-wise square root.

square(a, /, *[, stream])

Element-wise square.

squeeze(a, /[, axis, stream])

Remove length one axes from an array.

stack(arrays[, axis, stream])

Stacks the arrays along a new axis.

std(a, /[, axis, keepdims, ddof, stream])

Compute the standard deviation(s) over the given axes.

stop_gradient(a, /, *[, stream])

Stop gradients from being computed.

subtract(a, b[, stream])

Element-wise subtraction.

sum(a, /[, axis, keepdims, stream])

Sum reduce the array over the given axes.

swapaxes(a, /, axis1, axis2, *[, stream])

Swap two axes of an array.

take(a, /, indices[, axis, stream])

Take elements along an axis.

take_along_axis(a, /, indices[, axis, stream])

Take values along an axis at the specified indices.

tan(a, /, *[, stream])

Element-wise tangent.

tanh(a, /, *[, stream])

Element-wise hyperbolic tangent.

tensordot(a, b, /[, axes, stream])

Compute the tensor dot product along the specified axes.

tile(a, reps, /, *[, stream])

Construct an array by repeating a the number of times given by reps.

topk(a, /, k[, axis, stream])

Returns the k largest elements from the input along a given axis.

transpose(a, /[, axes, stream])

Transpose the dimensions of the array.

tri(n, m, k[, dtype, stream])

An array with ones at and below the given diagonal and zeros elsewhere.

tril(x, k, *[, stream])

Zeros the array above the given diagonal.

triu(x, k, *[, stream])

Zeros the array below the given diagonal.

var(a, /[, axis, keepdims, ddof, stream])

Compute the variance(s) over the given axes.

where(condition, x, y, /, *[, stream])

Select from x or y according to condition.

zeros(shape[, dtype, stream])

Construct an array of zeros.

zeros_like(a, /, *[, stream])

An array of zeros like the input.