Expand description
Kernel module
Modules§
- attention
- Attention kernels
- conv
- Convolution kernels
- grid_
sample - Grid sampling kernels
- interpolate
- Interpolation kernels
- matmul
- Matmul kernels
- pool
- Pooling kernels
- prng
- Pseudo-random number generator kernels
- quantization
- Quantization operations
- reduce
- Reduction algorithms
Enums§
- Mask
Fill Strategy - Define how to run the mask fill kernel.
- Mask
Where Strategy - Define how to run the mask where kernel.
Traits§
- Kernel
Metadata - Implement this trait to create a kernel definition.
Functions§
- bool_
cast - Cast a bool tensor to the given element type.
- cast
- Cast a tensor to the given element type.
- into_
contiguous - Make a jit tensor contiguous.
- into_
contiguous_ aligned - Make a jit tensor contiguous with an aligned last stride. Tensor is considered already contiguous if runtime can read it as is. This is equivalent in practice.
- mask_
fill - Execute the mask fill kernel with the given strategy.
- mask_
where - Execute the mask where kernel with the given strategy.
- slice
- Slice a jit tensor with a set of ranges
- slice_
with_ steps - Slice a tensor with steps