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, stream])

Approximate comparison of two arrays.

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

An or reduction over the given axes.

arange(start, stop, step[, dtype, stream])

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.

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.

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.

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

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

Element-wise cosine.

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

Element-wise hyperbolic cosine.

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

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

divide(a, b[, stream])

Element-wise division.

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.

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.

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, stream])

Load array(s) from a binary file in .npy, .npz, or .safetensors format.

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.

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])

An 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.

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

An 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.

ones(shape[, dtype, stream])

Construct an array of ones.

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

An array of ones like the input.

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.

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.

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

Round to the given number of decimals.

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

Element-wise reciprocal and square root.

save(file, arr, /[, retain_graph])

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_safetensors(file, arrays, /[, retain_graph])

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.

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, /[, dims, stream])

Compute the tensor dot product along the specified axes.

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.