Module ops

Traits

ArangeArgs

Represents the arguments for the arange function.

ArgReduceArgs

AvgPool1dArgs

AvgPool2dArgs

ClampBound

FlashAttentionOpts

FlattenArgs

MaxPool1dArgs

MaxPool2dArgs

ReduceArgs

ToPair

VarArgs

Functions

abs

add

Adds two Tensor objects.

arange

Creates a 1-D Tensor with values from a range.

argmax

argmin

avg_pool1d

avg_pool2d

broadcast_left

broadcast_right

broadcast_to

broadcast_to_unchecked

cat

chunk

clamp

conv1d

conv2d

conv_transpose1d

conv_transpose2d

cos

div

Divides two Tensor objects.

dropout

eq

erf

exp

flash_attention

flatten

from_array

Creates a Tensor from an array of values.

from_safetensor

Creates a Tensor from a safetensors::TensorView.

full

Creates a Tensor filled with a specified value.

full_like

Creates a Tensor filled with a specified value, with the same shape, data type and device as another Tensor.

gather

ge

gelu

gt

index_add

index_select

le

linspace

log

log2

log10

log_softmax

lt

matmul

max

max_pool1d

max_pool2d

maximum

mean

min

minimum

mul

Multiplies two Tensor objects.

narrow

ne

neg

new_gelu

ones

Creates a Tensor filled with ones.

ones_like

Creates a Tensor filled with ones, with the same shape, data type and device as another Tensor.

permute

powf

rand

Creates a Tensor filled with random values from a uniform distribution on the interval [0, 1).

rand_like

rand_with

randn

Creates a Tensor filled with random values from a normal distribution with mean 0 and variance 1.

randn_like

randn_with

relu

relu2

relu6

reshape

rsqrt

scatter_add

sign

silu

sin

softmax

sqrt

square

squeeze

sub

Subtracts two Tensor objects.

sum

tanh

to_contiguous

to_device

to_dtype

transpose

unsqueeze

upsample_nearest1d

upsample_nearest2d

var

where_cond

zeros

Creates a Tensor filled with zeros.

zeros_like

Creates a Tensor filled with zeros, with the same shape, data type and device as another Tensor.