Struct Flex 

pub struct Flex {}

The Flex backend, a fast, portable CPU backend for Burn.

The E and I type parameters exist purely to match the shape of other Burn backends (e.g. NdArray<E, I, Q>) so Flex slots into burn-dispatch’s generic dispatch macros. The body of Flex uses runtime DType dispatch, so both parameters are phantom and unused at runtime.

Limitations of the phantom generics

The Backend impl is provided only for the default instantiation Flex<f32, i32>. Writing Flex (with no arguments) resolves to the default and works exactly as before. Writing Flex<f64, i64> or any other non-default combination is a valid Rust type but will not satisfy trait bounds requiring Backend, producing errors like:

the trait bound `Flex<f64, i64>: Backend` is not satisfied

This is a deliberate compromise for the initial migration: making Flex generic over element types at the trait-impl level is a follow-up that would require rewriting all impl FooOps<Flex> for Flex blocks plus internal Flex::method() calls (tracked in #4762). Until then, treat the generic parameters as opaque shape placeholders; real element-type selection happens at runtime via DType.

The bound is locked in by a compile-fail doctest so that if someone later makes the Backend impl generic over E/I, this documentation gets flagged as out of date:

use burn_backend::Backend;
use burn_flex::Flex;
fn requires_backend<B: Backend>() {}
requires_backend::<Flex<f64, i64>>();

Trait Implementations

impl ActivationOps<Flex> for Flex

fn relu(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Applies the ReLU activation function.

fn relu_backward( output: FloatTensor<Flex>, grad: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Applies the ReLU activation function backward.

fn leaky_relu( tensor: FloatTensor<Flex>, negative_slope: Scalar, ) -> FloatTensor<Flex>

Applies the LeakyReLU activation function.

fn prelu( tensor: FloatTensor<Flex>, alpha: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Applies the PReLu activation function.

fn gelu(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Applies the Gelu activation function.

fn gelu_backward( x: FloatTensor<Flex>, grad: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Applies the Gelu activation function backward.

fn sigmoid(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Applies the Sigmoid activation function.

fn sigmoid_backward( output: FloatTensor<Flex>, grad: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Applies the Sigmoid activation function backward.

fn hard_sigmoid( tensor: FloatTensor<Flex>, alpha: Scalar, beta: Scalar, ) -> FloatTensor<Flex>

Applies the hard Sigmoid activation function.

fn log_sigmoid(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Applies the LogSigmoid activation function.

fn log_sigmoid_backward( x: FloatTensor<Flex>, grad: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Applies the LogSigmoid activation function backward.

fn softmax(tensor: FloatTensor<Flex>, dim: usize) -> FloatTensor<Flex>

Applies the softmax function along the given dimension.

fn log_softmax( tensor: <B as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <B as BackendTypes>::FloatTensorPrimitive

Applies the log-softmax function along the given dimension.

fn softmin( tensor: <B as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <B as BackendTypes>::FloatTensorPrimitive

Applies the softmin function along the given dimension.

impl Backend for Flex

fn name(_device: &Self::Device) -> String

Name of the backend.

fn seed(_device: &Self::Device, seed: u64)

Seeds the backend on the specified device.

fn device_count(_type_id: u16) -> usize

Returns the number of devices available on this backend. device is a reference device used to determine the underlying backend that should be queried. A CUDA device will return all devices available to CUDA, a Vulkan device will return all devices available to Vulkan, etc.

fn dtype_usage(_device: &Self::Device, dtype: DType) -> DTypeUsageSet

Returns the DTypeUsageSet for the given DType on the specified device.

fn flush(_device: &Self::Device)

Flush any pending operation of the backend.

fn ad_enabled(_device: &Self::Device) -> bool

If autodiff is enabled.

fn memory_persistent_allocations<Output, Input, Func>( device: &Self::Device, input: Input, func: Func, ) -> Output

where Output: Send, Input: Send, Func: Fn(Input) -> Output + Send,

Sets the current allocation mode to persistent.

fn memory_cleanup(device: &Self::Device)

Manually triggers a memory cleanup on the given device.

fn sync(_device: &Self::Device) -> Result<(), ExecutionError>

Sync the backend, ensure that all computation are finished.

fn graph_prepare(_device: &Self::Device) -> Result<(), ExecutionError>

Prepare device for an upcoming graph capture: route allocations into a stable pool so every buffer allocated before graph_stop_capture can be pinned. Call before the warmup run. No-op by default.

fn graph_start_capture(_device: &Self::Device) -> Result<(), ExecutionError>

Begin recording launches on device into a graph (see graph_stop_capture). Errors on backends without hardware graph support, so callers fall back to re-running.

fn graph_stop_capture( _device: &Self::Device, ) -> Result<Self::GraphPrimitive, ExecutionError>

Stop recording and return the captured graph, ready to graph_replay.

unsafe fn graph_replay( _device: &Self::Device, _graph: &Self::GraphPrimitive, ) -> Result<(), ExecutionError>

Replay a captured graph — one dispatch re-running the recorded launches against their original buffers.

fn staging<'a, Iter>(_data: Iter, _device: &Self::Device)

where Iter: Iterator<Item = &'a mut TensorData>,

Marks the given data as being used as a staging buffer for transfer between CPU and accelerators like GPUs.

fn supports_dtype(device: &Self::Device, dtype: DType) -> bool

Whether the type is fully supported by the specified device for general operations.

impl BackendIr for Flex

type Handle = HandleKind<Flex>

The type that can be used to point to a tensor of any kind.

fn float_tensor(handle: TensorHandle<Self::Handle>) -> FlexTensor

Convert a handle to a float tensor.

fn int_tensor(handle: TensorHandle<Self::Handle>) -> FlexTensor

Convert a handle to an int tensor.

fn bool_tensor(handle: TensorHandle<Self::Handle>) -> FlexTensor

Convert a handle to a bool tensor.

fn quantized_tensor(handle: TensorHandle<Self::Handle>) -> FlexQTensor

Convert a handle to a quantized tensor.

fn float_tensor_handle(tensor: FlexTensor) -> Self::Handle

Convert a float tensor to a handle.

fn int_tensor_handle(tensor: FlexTensor) -> Self::Handle

Convert an int tensor to a handle.

fn bool_tensor_handle(tensor: FlexTensor) -> Self::Handle

Convert a bool tensor to a handle.

fn quantized_tensor_handle(tensor: FlexQTensor) -> Self::Handle

Convert a quantized tensor to a handle.

impl BackendTypes for Flex

type Device = FlexDevice

Device type.

type FloatTensorPrimitive = FlexTensor

Tensor primitive to be used for all float operations.

type IntTensorPrimitive = FlexTensor

Tensor primitive to be used for all int operations.

type BoolTensorPrimitive = FlexTensor

Tensor primitive to be used for all bool operations.

type QuantizedTensorPrimitive = FlexQTensor

Tensor primitive to be used for all quantized operations.

type GraphPrimitive = GraphUnsupported

Captured graph primitive returned by Backend::graph_stop_capture and consumed by Backend::graph_replay: a backend-owned recording of a launch sequence that replays as a single dispatch.

impl BoolTensorOps<Flex> for Flex

fn bool_from_data(data: TensorData, _device: &Device<Flex>) -> BoolTensor<Flex>

Creates a tensor from the data structure.

async fn bool_into_data( tensor: BoolTensor<Flex>, ) -> Result<TensorData, ExecutionError>

Converts the tensor to a data structure.

fn bool_to_device( tensor: BoolTensor<Flex>, _device: &Device<Flex>, ) -> BoolTensor<Flex>

Moves the tensor to the device.

fn bool_cat(tensors: Vec<BoolTensor<Flex>>, dim: usize) -> BoolTensor<Flex>

Concatenates the tensors along the given dimension.

fn bool_reshape(tensor: BoolTensor<Flex>, shape: Shape) -> BoolTensor<Flex>

Reshapes the tensor.

fn bool_slice(tensor: BoolTensor<Flex>, slices: &[Slice]) -> BoolTensor<Flex>

Gets the values from the tensor for the given ranges.

fn bool_empty( shape: Shape, _device: &Device<Flex>, dtype: BoolDType, ) -> BoolTensor<Flex>

Creates a new bool tensor.

fn bool_slice_assign( tensor: BoolTensor<Flex>, slices: &[Slice], value: BoolTensor<Flex>, ) -> BoolTensor<Flex>

Sets the values in the tensor for the given ranges.

fn bool_into_int( tensor: BoolTensor<Flex>, out_dtype: IntDType, ) -> IntTensor<Flex>

Converts bool tensor to int tensor.

fn bool_into_float( tensor: BoolTensor<Flex>, out_dtype: FloatDType, ) -> FloatTensor<Flex>

Converts bool tensor to float tensor.

fn bool_swap_dims( tensor: BoolTensor<Flex>, dim1: usize, dim2: usize, ) -> BoolTensor<Flex>

Swaps two dimensions of a bool tensor.

fn bool_permute(tensor: BoolTensor<Flex>, axes: &[usize]) -> BoolTensor<Flex>

Permutes the dimensions of a tensor.

fn bool_flip(tensor: BoolTensor<Flex>, axes: &[usize]) -> BoolTensor<Flex>

Reverse the order of elements in a tensor along the given axes.

fn bool_equal(lhs: BoolTensor<Flex>, rhs: BoolTensor<Flex>) -> BoolTensor<Flex>

Equates the two tensors.

fn bool_not(tensor: BoolTensor<Flex>) -> BoolTensor<Flex>

Inverses boolean values.

fn bool_and(lhs: BoolTensor<Flex>, rhs: BoolTensor<Flex>) -> BoolTensor<Flex>

Executes the logical and (&&) operation on two boolean tensors.

fn bool_or(lhs: BoolTensor<Flex>, rhs: BoolTensor<Flex>) -> BoolTensor<Flex>

Executes the logical or (||) operation on two boolean tensors.

fn bool_xor(lhs: BoolTensor<Flex>, rhs: BoolTensor<Flex>) -> BoolTensor<Flex>

Element-wise exclusive or.

fn bool_expand(tensor: BoolTensor<Flex>, shape: Shape) -> BoolTensor<Flex>

Broadcasts the bool tensor to the given shape.

fn bool_zeros( shape: Shape, device: &Device<Flex>, dtype: BoolDType, ) -> BoolTensor<Flex>

Creates a new bool tensor filled false.

fn bool_ones( shape: Shape, _device: &Device<Flex>, dtype: BoolDType, ) -> BoolTensor<Flex>

Creates a new bool tensor filled true.

fn bool_mask_where( tensor: BoolTensor<Flex>, mask: BoolTensor<Flex>, value: BoolTensor<Flex>, ) -> BoolTensor<Flex>

Fills the tensor with values from the value tensor if the mask is true at the given indices.

fn bool_mask_fill( tensor: BoolTensor<Flex>, mask: BoolTensor<Flex>, value: Scalar, ) -> BoolTensor<Flex>

Fills the tensor with the given value if the mask is true at the given indices.

fn bool_gather( dim: usize, tensor: BoolTensor<Flex>, indices: IntTensor<Flex>, ) -> BoolTensor<Flex>

Gather elements from the tensor at the given indices.

fn bool_scatter_or( dim: usize, tensor: BoolTensor<Flex>, indices: IntTensor<Flex>, value: BoolTensor<Flex>, ) -> BoolTensor<Flex>

Scatter a given value to the tensor at the given indices using boolean or reduction.

fn bool_equal_elem(lhs: BoolTensor<Flex>, rhs: Scalar) -> BoolTensor<Flex>

Element-wise equality comparison with a scalar.

fn bool_unfold( tensor: BoolTensor<Flex>, dim: usize, size: usize, step: usize, ) -> BoolTensor<Flex>

Unfold windows along a dimension.

fn bool_not_equal( lhs: BoolTensor<Flex>, rhs: BoolTensor<Flex>, ) -> BoolTensor<Flex>

Element-wise non-equality comparison.

fn bool_not_equal_elem(lhs: BoolTensor<Flex>, rhs: Scalar) -> BoolTensor<Flex>

Element-wise non-equality comparison with a scalar.

fn bool_any(tensor: BoolTensor<Flex>) -> BoolTensor<Flex>

Tests if any element in the boolean tensor evaluates to True.

fn bool_any_dim(tensor: BoolTensor<Flex>, dim: usize) -> BoolTensor<Flex>

Tests if any element in the boolean tensor evaluates to True along a given dimension dim.

fn bool_all(tensor: BoolTensor<Flex>) -> BoolTensor<Flex>

Tests if all elements in the boolean tensor evaluate to True.

fn bool_all_dim(tensor: BoolTensor<Flex>, dim: usize) -> BoolTensor<Flex>

Tests if all elements in the boolean tensor evaluate to True along a given dimension dim.

fn bool_select( tensor: BoolTensor<Flex>, dim: usize, indices: IntTensor<Flex>, ) -> BoolTensor<Flex>

Select tensor elements along the given dimension corresponding to the given indices.

fn bool_select_or( tensor: BoolTensor<Flex>, dim: usize, indices: IntTensor<Flex>, value: BoolTensor<Flex>, ) -> BoolTensor<Flex>

Assign the selected elements along the given dimension corresponding to the given indices to the given value using sum reduction.

fn bool_transpose(tensor: BoolTensor<Flex>) -> BoolTensor<Flex>

Transposes a bool tensor.

fn bool_repeat_dim( tensor: BoolTensor<Flex>, dim: usize, times: usize, ) -> BoolTensor<Flex>

Repeats one dimension of the tensor a given number of times along that dimension.

async fn bool_argwhere( tensor: BoolTensor<Flex>, out_dtype: IntDType, ) -> IntTensor<Flex>

Compute the indices of the elements that are non-zero, grouped by element.

fn bool_mask_select( tensor: <B as BackendTypes>::BoolTensorPrimitive, mask: <B as BackendTypes>::BoolTensorPrimitive, ) -> impl Future<Output = <B as BackendTypes>::BoolTensorPrimitive> + Send + 'static

Selects the elements of the tensor where the mask is true, returned as a 1D tensor.

impl Clone for Flex

fn clone(&self) -> Flex

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Copy for Flex

impl Debug for Flex

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for Flex

fn default() -> Flex

Returns the “default value” for a type.

impl DistributedOps<Flex> for Flex

fn start_communication_server( devices: &[<B as BackendTypes>::Device], config: DistributedConfig, )

Start the communication server used to orchestrate tensor syncing between devices.

fn close_communication_server(_device: &<B as BackendTypes>::Device)

Close the communication server used to orchestrate syncing between devices.

fn register_sync_parameters( _device: &<B as BackendTypes>::Device, distributed_params: Vec<DistributedParams>, )

Register the parameters that will require gradient synchronization for the upcoming backward pass.

fn submit_sync_collective(device: &<B as BackendTypes>::Device)

Tell the gradient sync server that this device has submitted all its sync operations and is ready to be synchronized.

fn submit_gradient_sync( tensor: TensorRef<B>, distributed_params: DistributedParams, )

Submit a gradient tensor for synchronization across all devices.

fn all_reduce( _tensor: <B as BackendTypes>::FloatTensorPrimitive, _op: ReduceOperation, _device_ids: Vec<DeviceId>, ) -> CollectiveTensor<B>

all_reduce operation.

fn sync_collective(_device: &<B as BackendTypes>::Device)

Sync the collective operations.

unsafe fn comm_device(tensor: &TensorRef<B>) -> <B as BackendTypes>::Device

Get the device of the tensor reference.

unsafe fn float_from_ref( tensor: &TensorRef<B>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Returns a clone of the float tensor from the tensor reference.

impl FloatTensorOps<Flex> for Flex

fn float_from_data( data: TensorData, _device: &Device<Flex>, ) -> FloatTensor<Flex>

Creates a new tensor from the data structure.

fn float_random( shape: Shape, distribution: Distribution, _device: &Device<Flex>, dtype: FloatDType, ) -> FloatTensor<Flex>

Creates a new tensor with random values.

async fn float_into_data( tensor: FloatTensor<Flex>, ) -> Result<TensorData, ExecutionError>

Converts the tensor to a data structure.

fn float_to_device( tensor: FloatTensor<Flex>, _device: &Device<Flex>, ) -> FloatTensor<Flex>

Moves the tensor to the given device.

fn float_detach(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Detaches a tensor from the computation graph.

fn float_into_int( tensor: FloatTensor<Flex>, out_dtype: IntDType, ) -> IntTensor<Flex>

Converts float tensor to int tensor.

fn float_empty( shape: Shape, _device: &Device<Flex>, dtype: FloatDType, ) -> FloatTensor<Flex>

Creates an empty tensor with the given shape.

fn float_add( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Adds two tensors together.

fn float_add_scalar(lhs: FloatTensor<Flex>, rhs: Scalar) -> FloatTensor<Flex>

Adds a scalar to a tensor.

fn float_sub( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Subtracts two tensors.

fn float_sub_scalar(lhs: FloatTensor<Flex>, rhs: Scalar) -> FloatTensor<Flex>

Subtracts a scalar from a tensor.

fn float_mul( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Multiplies two tensors together element-wise.

fn float_mul_scalar(lhs: FloatTensor<Flex>, rhs: Scalar) -> FloatTensor<Flex>

Multiplies a tensor by a scalar.

fn float_div( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Divides two tensors element-wise.

fn float_div_scalar(lhs: FloatTensor<Flex>, rhs: Scalar) -> FloatTensor<Flex>

Divides a tensor by a scalar.

fn float_remainder( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Computes the remainder of division between two tensors element-wise.

fn float_remainder_scalar( lhs: FloatTensor<Flex>, rhs: Scalar, ) -> FloatTensor<Flex>

Computes the modulus of a tensor given a scalar.

fn float_matmul( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Multiplies two tensors together using matrix multiplication.

fn float_cross( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, dim: usize, ) -> FloatTensor<Flex>

Computes the cross product of two tensors along a given dimension.

fn float_recip(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Calculates the reciprocals element-wise

fn float_swap_dims( tensor: FloatTensor<Flex>, dim1: usize, dim2: usize, ) -> FloatTensor<Flex>

Swaps two dimensions of a tensor.

fn float_permute(tensor: FloatTensor<Flex>, axes: &[usize]) -> FloatTensor<Flex>

Permutes the dimensions of a tensor.

fn float_flip(tensor: FloatTensor<Flex>, axes: &[usize]) -> FloatTensor<Flex>

Reverse the order of elements in a tensor along the given axes.

fn float_cat(tensors: Vec<FloatTensor<Flex>>, dim: usize) -> FloatTensor<Flex>

Concatenates tensors along a dimension.

fn float_reshape(tensor: FloatTensor<Flex>, shape: Shape) -> FloatTensor<Flex>

Reshapes a tensor.

fn float_gather( dim: usize, tensor: FloatTensor<Flex>, indices: IntTensor<Flex>, ) -> FloatTensor<Flex>

Gather elements from a tensor.

fn float_scatter_add( dim: usize, tensor: FloatTensor<Flex>, indices: IntTensor<Flex>, value: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Scatter elements into a tensor using sum reduction.

fn float_scatter_nd( data: FloatTensor<Flex>, indices: IntTensor<Flex>, values: FloatTensor<Flex>, reduction: IndexingUpdateOp, ) -> FloatTensor<Flex>

Multi-dimensional scatter: update data at locations specified by indices with values.

fn float_gather_nd( data: FloatTensor<Flex>, indices: IntTensor<Flex>, ) -> FloatTensor<Flex>

Multi-dimensional gather: collect slices from data at locations specified by indices.

fn float_select( tensor: FloatTensor<Flex>, dim: usize, indices: IntTensor<Flex>, ) -> FloatTensor<Flex>

Select tensor elements along the given dimension corresponding for the given indices.

fn float_select_add( tensor: FloatTensor<Flex>, dim: usize, indices: IntTensor<Flex>, value: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Assign the selected elements along the given dimension corresponding for the given indices to the given value using sum reduction.

fn float_slice(tensor: FloatTensor<Flex>, slices: &[Slice]) -> FloatTensor<Flex>

Select tensor elements corresponding to the given slices.

fn float_slice_assign( tensor: FloatTensor<Flex>, slices: &[Slice], value: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Assign the selected elements corresponding to the given slices to the given value.

fn float_mask_where( tensor: FloatTensor<Flex>, mask: BoolTensor<Flex>, value: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Update the given tensor with the value tensor where the mask is true.

fn float_mask_fill( tensor: FloatTensor<Flex>, mask: BoolTensor<Flex>, value: Scalar, ) -> FloatTensor<Flex>

Update the given tensor with the value where the mask is true.

fn float_equal( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Equal comparison of two tensors.

fn float_equal_elem( lhs: FloatTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Equal comparison of a tensor and a scalar.

fn float_greater( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Greater than comparison of two tensors.

fn float_greater_elem( lhs: FloatTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Greater than comparison of a tensor and a scalar.

fn float_greater_equal( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Greater than or equal comparison of two tensors.

fn float_greater_equal_elem( lhs: FloatTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Greater than or equal comparison of a tensor and a scalar.

fn float_lower( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Less than comparison of two tensors.

fn float_lower_elem( lhs: FloatTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Less than comparison of a tensor and a scalar.

fn float_lower_equal( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Less than or equal comparison of two tensors.

fn float_lower_equal_elem( lhs: FloatTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Less than or equal comparison of a tensor and a scalar.

fn float_not_equal( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise non-equality comparison.

fn float_not_equal_elem( lhs: FloatTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise non-equality comparison with a scalar.

fn float_neg(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Negates a tensor element-wise.

fn float_clamp( tensor: FloatTensor<Flex>, min: Scalar, max: Scalar, ) -> FloatTensor<Flex>

Clamps a tensor between a minimum and maximum value.

fn float_clamp_min(tensor: FloatTensor<Flex>, min: Scalar) -> FloatTensor<Flex>

Clamps a tensor under a minimum value.

fn float_clamp_max(tensor: FloatTensor<Flex>, max: Scalar) -> FloatTensor<Flex>

Clamps a tensor over a maximum value.

fn float_sign(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns the signs of the float tensor.

fn float_mean(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Mean of all elements in a tensor.

fn float_max(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Gets the maximum element of a tensor.

fn float_max_dim(tensor: FloatTensor<Flex>, dim: usize) -> FloatTensor<Flex>

Gets the maximum elements of a tensor along an axis.

fn float_min(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Gets the minimum element of a tensor.

fn float_min_dim(tensor: FloatTensor<Flex>, dim: usize) -> FloatTensor<Flex>

Gets the minimum elements of a tensor along an axis.

fn float_max_dim_with_indices( tensor: FloatTensor<Flex>, dim: usize, indices_dtype: IntDType, ) -> (FloatTensor<Flex>, IntTensor<Flex>)

Gets the maximum elements of a tensor along an axis and their indices.

fn float_min_dim_with_indices( tensor: FloatTensor<Flex>, dim: usize, indices_dtype: IntDType, ) -> (FloatTensor<Flex>, IntTensor<Flex>)

Gets the minimum elements of a tensor along an axis and their indices.

fn float_any( tensor: FloatTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Tests if any element in the float tensor evaluates to True.

fn float_any_dim( tensor: FloatTensor<Flex>, dim: usize, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Tests if any element in the float tensor evaluates to True along a given dimension dim.

fn float_all( tensor: FloatTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Tests if all elements in the float tensor evaluate to True.

fn float_all_dim( tensor: FloatTensor<Flex>, dim: usize, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Tests if all elements in the float tensor evaluate to True along a given dimension dim.

fn float_sum(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Sum of all elements in a tensor.

fn float_sum_dim(tensor: FloatTensor<Flex>, dim: usize) -> FloatTensor<Flex>

Sum of all elements in a tensor along a dimension.

fn float_mean_dim(tensor: FloatTensor<Flex>, dim: usize) -> FloatTensor<Flex>

Mean of all elements in a tensor along a dimension.

fn float_prod(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Product of all elements in a tensor.

fn float_prod_dim(tensor: FloatTensor<Flex>, dim: usize) -> FloatTensor<Flex>

Product of all elements in a tensor along a dimension.

fn float_cumsum(tensor: FloatTensor<Flex>, dim: usize) -> FloatTensor<Flex>

Computes the cumulative sum of elements along a dimension.

fn float_cumprod(tensor: FloatTensor<Flex>, dim: usize) -> FloatTensor<Flex>

Computes the cumulative product of elements along a dimension.

fn float_cummin(tensor: FloatTensor<Flex>, dim: usize) -> FloatTensor<Flex>

Computes the cumulative minimum of elements along a dimension.

fn float_cummax(tensor: FloatTensor<Flex>, dim: usize) -> FloatTensor<Flex>

Computes the cumulative maximum of elements along a dimension.

fn float_cast(tensor: FloatTensor<Flex>, dtype: FloatDType) -> FloatTensor<Flex>

Converts a tensor to another floating point data type.

fn float_exp(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with exponential values.

fn float_log(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with natural logarithm values.

fn float_log1p(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with logarithm values of (1 + Xi).

fn float_powf( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Element-wise power with a FloatTensor.

fn float_powf_scalar_impl( tensor: FloatTensor<Flex>, value: Scalar, ) -> FloatTensor<Flex>

Returns a new tensor with values raised to the power of float value.

fn float_sqrt(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with square root values.

fn float_abs(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with absolute values.

fn float_cos(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with cosine values.

fn float_sin(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with sine values.

fn float_tan(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with tangent values.

fn float_cosh(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with hyperbolic cosine values.

fn float_sinh(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with hyperbolic sine values.

fn float_tanh(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with hyperbolic tangent values.

fn float_acos(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with inverse cosine values.

fn float_acosh(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with inverse hyperbolic cosine values.

fn float_asin(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with inverse sine values.

fn float_asinh(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with inverse hyperbolic sine values.

fn float_atan(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with the inverse tangent values.

fn float_atanh(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with the inverse hyperbolic tangent values.

fn float_atan2( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Returns a tensor with the four-quadrant inverse tangent values of y and x.

fn float_round(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with rounded values.

fn float_floor(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with floored values.

fn float_ceil(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with ceiled values.

fn float_trunc(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with truncated values.

fn float_erf(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Returns a new tensor with the error function values.

fn float_argmax( tensor: FloatTensor<Flex>, dim: usize, out_dtype: IntDType, ) -> IntTensor<Flex>

Gets the indices of the maximum elements of a tensor along an axis.

fn float_argmin( tensor: FloatTensor<Flex>, dim: usize, out_dtype: IntDType, ) -> IntTensor<Flex>

Gets the indices of the minimum elements of a tensor along an axis.

fn float_expand(tensor: FloatTensor<Flex>, shape: Shape) -> FloatTensor<Flex>

Broadcasts the float tensor to the given shape.

fn float_unfold( tensor: FloatTensor<Flex>, dim: usize, size: usize, step: usize, ) -> FloatTensor<Flex>

Unfold windows along a dimension.

fn float_grid_sample_2d( tensor: FloatTensor<Flex>, grid: FloatTensor<Flex>, options: GridSampleOptions, ) -> FloatTensor<Flex>

Samples tensor as a two-dimensional spatial grid of (possibly multi-channel) values, using the given locations in [-1, 1].

fn float_zeros( shape: Shape, _device: &Device<Flex>, dtype: FloatDType, ) -> FloatTensor<Flex>

Creates a new tensor with zeros.

fn float_ones( shape: Shape, _device: &Device<Flex>, dtype: FloatDType, ) -> FloatTensor<Flex>

Creates a new tensor with ones.

fn float_full( shape: Shape, fill_value: Scalar, _device: &Device<Flex>, dtype: FloatDType, ) -> FloatTensor<Flex>

Creates a tensor filled with given value.

fn float_transpose(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Transposes a tensor.

fn float_repeat_dim( tensor: FloatTensor<Flex>, dim: usize, times: usize, ) -> FloatTensor<Flex>

Repeat the tensor along the given dimension.

fn float_sort( tensor: FloatTensor<Flex>, dim: usize, descending: bool, ) -> FloatTensor<Flex>

Sort the elements of the input tensor by value in along a given dimension.

fn float_sort_with_indices( tensor: FloatTensor<Flex>, dim: usize, descending: bool, indices_dtype: IntDType, ) -> (FloatTensor<Flex>, IntTensor<Flex>)

Sort the elements of the input tensor by value in along a given dimension.

fn float_argsort( tensor: FloatTensor<Flex>, dim: usize, descending: bool, out_dtype: IntDType, ) -> IntTensor<Flex>

Returns the indices that sort the elements of the input tensor by value along a given dimension.

fn float_powi(lhs: FloatTensor<Flex>, rhs: IntTensor<Flex>) -> FloatTensor<Flex>

Element-wise power with an IntTensor.

fn float_powi_scalar(lhs: FloatTensor<Flex>, rhs: Scalar) -> FloatTensor<Flex>

Raises a tensor to the power of an int scalar.

fn float_powf_scalar( tensor: FloatTensor<Flex>, value: Scalar, ) -> FloatTensor<Flex>

Returns a new tensor with values raised to the power of float value.

fn float_max_abs(tensor: FloatTensor<Flex>) -> FloatTensor<Flex>

Gets the maximum absolute element of a tensor.

fn float_max_abs_dim(tensor: FloatTensor<Flex>, dim: usize) -> FloatTensor<Flex>

Gets the maximum absolute elements of a tensor along an axis.

fn float_is_nan( tensor: FloatTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Returns a new tensor with boolean elements indicating whether each element of the input is NaN.

fn float_is_inf( tensor: FloatTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Returns a new tensor with boolean elements indicating whether each element of the input is infinite (either +INF or -INF).

fn float_hypot( lhs: FloatTensor<Flex>, rhs: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Returns a new tensor with the Euclidean distance values.

fn float_scatter( dim: usize, tensor: <B as BackendTypes>::FloatTensorPrimitive, indices: <B as BackendTypes>::IntTensorPrimitive, value: <B as BackendTypes>::FloatTensorPrimitive, update: IndexingUpdateOp, ) -> <B as BackendTypes>::FloatTensorPrimitive

Scatter elements into a tensor using the specified update operation.

fn float_select_assign( tensor: <B as BackendTypes>::FloatTensorPrimitive, dim: usize, indices: <B as BackendTypes>::IntTensorPrimitive, value: <B as BackendTypes>::FloatTensorPrimitive, update: IndexingUpdateOp, ) -> <B as BackendTypes>::FloatTensorPrimitive

Assign selected elements along a dimension using the specified update operation.

fn float_mask_select( tensor: <B as BackendTypes>::FloatTensorPrimitive, mask: <B as BackendTypes>::BoolTensorPrimitive, ) -> impl Future<Output = <B as BackendTypes>::FloatTensorPrimitive> + Send + 'static

Selects the elements of the tensor where the mask is true, returned as a 1D tensor.

fn float_set_require_grad( tensor: <B as BackendTypes>::FloatTensorPrimitive, _require_grad: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Sets the require_grad flag of a tensor.

fn float_is_require_grad( _tensor: &<B as BackendTypes>::FloatTensorPrimitive, ) -> bool

Returns the require_grad flag of a tensor.

fn float_powi_scalar_impl( lhs: <B as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, ) -> <B as BackendTypes>::FloatTensorPrimitive

Raises a tensor to the power of an int scalar.

fn float_argtopk( tensor: <B as BackendTypes>::FloatTensorPrimitive, dim: usize, k: usize, out_dtype: IntDType, ) -> <B as BackendTypes>::IntTensorPrimitive

Gets the indices of the k maximum elements of a tensor along an axis. if two elements are equals, it will be ordered by lowest indices

fn float_topk( tensor: <B as BackendTypes>::FloatTensorPrimitive, dim: usize, k: usize, ) -> <B as BackendTypes>::FloatTensorPrimitive

Gets the values of the k maximum elements of a tensor along an axis.

fn float_topk_with_indices( tensor: <B as BackendTypes>::FloatTensorPrimitive, dim: usize, k: usize, out_dtype: IntDType, ) -> (<B as BackendTypes>::FloatTensorPrimitive, <B as BackendTypes>::IntTensorPrimitive)

Gets the values of the k maximum elements of a tensor along an axis, and their indices.

impl IntTensorOps<Flex> for Flex

fn int_gather( dim: usize, tensor: IntTensor<Flex>, indices: IntTensor<Flex>, ) -> IntTensor<Flex>

Gather ints along dim at the given indices.

The tensor dispatches on its own int dtype (I8/I16/I32/I64 signed or U8/U16/U32/U64 unsigned). The indices tensor may be any of those widths too - it’s normalised to isize by the shared read_indices helper in ops::gather_scatter before the kernel runs, so callers are not required to pre-convert to I64.

fn int_scatter_add( dim: usize, tensor: IntTensor<Flex>, indices: IntTensor<Flex>, value: IntTensor<Flex>, ) -> IntTensor<Flex>

Scatter-add int values at the given indices along dim.

tensor and value must share the same int dtype; indices may be any supported int width. See int_gather for the full index-width policy.

fn int_select( tensor: IntTensor<Flex>, dim: usize, indices: IntTensor<Flex>, ) -> IntTensor<Flex>

Select ints along dim by a 1D index tensor.

The indices tensor may be any supported int width. See int_gather for the full index-width policy.

fn int_select_add( tensor: IntTensor<Flex>, dim: usize, indices: IntTensor<Flex>, value: IntTensor<Flex>, ) -> IntTensor<Flex>

Select-add int values at a 1D index tensor along dim.

tensor and value must share the same int dtype; indices may be any supported int width. See int_gather for the full index-width policy.

fn int_from_data(data: TensorData, _device: &Device<Flex>) -> IntTensor<Flex>

Creates a tensor from the data structure.

async fn int_into_data( tensor: IntTensor<Flex>, ) -> Result<TensorData, ExecutionError>

Converts the tensor to a data structure.

fn int_to_device( tensor: IntTensor<Flex>, _device: &Device<Flex>, ) -> IntTensor<Flex>

Moves the tensor to the given device.

fn int_cat(tensors: Vec<IntTensor<Flex>>, dim: usize) -> IntTensor<Flex>

Concatenates the given tensors along the given dimension.

fn int_reshape(tensor: IntTensor<Flex>, shape: Shape) -> IntTensor<Flex>

Reshapes the tensor.

fn int_slice(tensor: IntTensor<Flex>, slices: &[Slice]) -> IntTensor<Flex>

Gets the element at the given indices.

fn int_empty( shape: Shape, _device: &Device<Flex>, dtype: IntDType, ) -> IntTensor<Flex>

Creates a new int tensor.

fn int_mask_where( tensor: IntTensor<Flex>, mask: BoolTensor<Flex>, value: IntTensor<Flex>, ) -> IntTensor<Flex>

Fills the tensor with values from the value tensor if the mask is true at the given indices.

fn int_mask_fill( tensor: IntTensor<Flex>, mask: BoolTensor<Flex>, value: Scalar, ) -> IntTensor<Flex>

Fills the tensor with the given value if the mask is true at the given indices.

fn int_slice_assign( tensor: IntTensor<Flex>, slices: &[Slice], value: IntTensor<Flex>, ) -> IntTensor<Flex>

Sets the values in the tensor for the given ranges.

fn int_scatter_nd( data: IntTensor<Flex>, indices: IntTensor<Flex>, values: IntTensor<Flex>, reduction: IndexingUpdateOp, ) -> IntTensor<Flex>

Multi-dimensional scatter for int tensors.

fn int_gather_nd( data: IntTensor<Flex>, indices: IntTensor<Flex>, ) -> IntTensor<Flex>

Multi-dimensional gather for int tensors.

fn int_equal( lhs: IntTensor<Flex>, rhs: IntTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise equality comparison.

fn int_equal_elem( lhs: IntTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise equality comparison with a scalar.

fn int_greater( lhs: IntTensor<Flex>, rhs: IntTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise greater than comparison.

fn int_greater_elem( lhs: IntTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise greater than comparison with a scalar.

fn int_greater_equal( lhs: IntTensor<Flex>, rhs: IntTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise greater than or equal comparison.

fn int_greater_equal_elem( lhs: IntTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise greater than or equal comparison with a scalar.

fn int_lower( lhs: IntTensor<Flex>, rhs: IntTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise less than comparison.

fn int_lower_elem( lhs: IntTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise less than comparison with a scalar.

fn int_lower_equal( lhs: IntTensor<Flex>, rhs: IntTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise less than or equal comparison.

fn int_lower_equal_elem( lhs: IntTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise less than or equal comparison with a scalar.

fn int_add(lhs: IntTensor<Flex>, rhs: IntTensor<Flex>) -> IntTensor<Flex>

Element-wise addition.

fn int_add_scalar(lhs: IntTensor<Flex>, rhs: Scalar) -> IntTensor<Flex>

Element-wise addition with a scalar.

fn int_sub(lhs: IntTensor<Flex>, rhs: IntTensor<Flex>) -> IntTensor<Flex>

Element-wise subtraction.

fn int_sub_scalar(lhs: IntTensor<Flex>, rhs: Scalar) -> IntTensor<Flex>

Element-wise subtraction with a scalar.

fn int_mul(lhs: IntTensor<Flex>, rhs: IntTensor<Flex>) -> IntTensor<Flex>

Element-wise multiplication.

fn int_mul_scalar(lhs: IntTensor<Flex>, rhs: Scalar) -> IntTensor<Flex>

Element-wise multiplication with a scalar.

fn int_div(lhs: IntTensor<Flex>, rhs: IntTensor<Flex>) -> IntTensor<Flex>

Element-wise division.

fn int_div_scalar(lhs: IntTensor<Flex>, rhs: Scalar) -> IntTensor<Flex>

Element-wise division with a scalar.

fn int_remainder(lhs: IntTensor<Flex>, rhs: IntTensor<Flex>) -> IntTensor<Flex>

Element-wise modulus.

fn int_remainder_scalar(lhs: IntTensor<Flex>, rhs: Scalar) -> IntTensor<Flex>

Element-wise modulus with a scalar.

fn int_into_float( tensor: IntTensor<Flex>, out_dtype: FloatDType, ) -> FloatTensor<Flex>

Converts int tensor to float tensor.

fn int_swap_dims( tensor: IntTensor<Flex>, dim1: usize, dim2: usize, ) -> IntTensor<Flex>

Swaps two dimensions of an int tensor.

fn int_permute(tensor: IntTensor<Flex>, axes: &[usize]) -> IntTensor<Flex>

Permutes the dimensions of a tensor.

fn int_flip(tensor: IntTensor<Flex>, axes: &[usize]) -> IntTensor<Flex>

Reverse the order of elements in a tensor along the given axes.

fn int_random( shape: Shape, distribution: Distribution, _device: &Device<Flex>, dtype: IntDType, ) -> IntTensor<Flex>

Creates a new int tensor with random values.

fn int_expand(tensor: IntTensor<Flex>, shape: Shape) -> IntTensor<Flex>

Broadcasts the int tensor to the given shape.

fn int_matmul(lhs: IntTensor<Flex>, rhs: IntTensor<Flex>) -> IntTensor<Flex>

Multiplies two tensors together using matrix multiplication.

fn int_sum(tensor: IntTensor<Flex>) -> IntTensor<Flex>

Sums all elements in the tensor.

fn int_sum_dim(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Sums all elements in the tensor along a dimension.

fn int_prod(tensor: IntTensor<Flex>) -> IntTensor<Flex>

Computes the product of all elements in the tensor.

fn int_prod_dim(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Computes the product of all elements in the tensor along a dimension.

fn int_mean_dim(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Computes the mean of all elements in the tensor along a dimension.

fn int_cumsum(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Computes the cumulative sum of elements along a dimension.

fn int_cumprod(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Computes the cumulative product of elements along a dimension.

fn int_cummin(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Computes the cumulative minimum of elements along a dimension.

fn int_cummax(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Computes the cumulative maximum of elements along a dimension.

fn int_argmax(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Gets the indices of the maximum elements along a dimension.

fn int_argmin(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Gets the indices of the minimum elements along a dimension.

fn int_abs(tensor: IntTensor<Flex>) -> IntTensor<Flex>

Returns a new tensor with absolute values.

fn bitwise_and(lhs: IntTensor<Flex>, rhs: IntTensor<Flex>) -> IntTensor<Flex>

Bitwise AND operation for Int Tensors

fn bitwise_and_scalar(lhs: IntTensor<Flex>, rhs: Scalar) -> IntTensor<Flex>

Bitwise AND operation for Int Tensors with a scalar

fn bitwise_or(lhs: IntTensor<Flex>, rhs: IntTensor<Flex>) -> IntTensor<Flex>

Bitwise OR operation for Int Tensors

fn bitwise_or_scalar(lhs: IntTensor<Flex>, rhs: Scalar) -> IntTensor<Flex>

Bitwise OR operation for Int Tensors with a scalar

fn bitwise_xor(lhs: IntTensor<Flex>, rhs: IntTensor<Flex>) -> IntTensor<Flex>

Bitwise XOR operation for Int Tensors

fn bitwise_xor_scalar(lhs: IntTensor<Flex>, rhs: Scalar) -> IntTensor<Flex>

Bitwise XOR operation for Int Tensors with a scalar

fn bitwise_not(tensor: IntTensor<Flex>) -> IntTensor<Flex>

Bitwise NOT operation for Int Tensors

fn bitwise_left_shift( lhs: IntTensor<Flex>, rhs: IntTensor<Flex>, ) -> IntTensor<Flex>

Bitwise left shift operation for Int Tensors

fn bitwise_left_shift_scalar( lhs: IntTensor<Flex>, rhs: Scalar, ) -> IntTensor<Flex>

Bitwise left shift operation for Int Tensors with a scalar

fn bitwise_right_shift( lhs: IntTensor<Flex>, rhs: IntTensor<Flex>, ) -> IntTensor<Flex>

Bitwise right shift operation for Int Tensors

fn bitwise_right_shift_scalar( lhs: IntTensor<Flex>, rhs: Scalar, ) -> IntTensor<Flex>

Bitwise right shift operation for Int Tensors with a scalar

fn int_cast(tensor: IntTensor<Flex>, dtype: IntDType) -> IntTensor<Flex>

Converts a tensor to another integer data type.

fn int_unfold( tensor: IntTensor<Flex>, dim: usize, size: usize, step: usize, ) -> IntTensor<Flex>

Unfold windows along a dimension.

fn int_neg(tensor: IntTensor<Flex>) -> IntTensor<Flex>

Element-wise negation.

fn int_clamp( tensor: IntTensor<Flex>, min: Scalar, max: Scalar, ) -> IntTensor<Flex>

Clamps a tensor between a minimum and maximum value.

fn int_clamp_min(tensor: IntTensor<Flex>, min: Scalar) -> IntTensor<Flex>

Clamps a tensor under a minimum value.

fn int_clamp_max(tensor: IntTensor<Flex>, max: Scalar) -> IntTensor<Flex>

Clamps a tensor over a maximum value.

fn int_sign(tensor: IntTensor<Flex>) -> IntTensor<Flex>

Returns the signs of the int tensor.

fn int_mean(tensor: IntTensor<Flex>) -> IntTensor<Flex>

Computes the mean of all elements in the tensor.

fn int_max(tensor: IntTensor<Flex>) -> IntTensor<Flex>

Gets the maximum element in the tensor.

fn int_max_dim(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Gets the maximum element in the tensor along a dimension.

fn int_min(tensor: IntTensor<Flex>) -> IntTensor<Flex>

Gets the minimum element in the tensor.

fn int_min_dim(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Gets the minimum elements in the tensor along a dimension.

fn int_max_dim_with_indices( tensor: IntTensor<Flex>, dim: usize, ) -> (IntTensor<Flex>, IntTensor<Flex>)

Gets the maximum elements and corresponding indices along a dimension.

fn int_min_dim_with_indices( tensor: IntTensor<Flex>, dim: usize, ) -> (IntTensor<Flex>, IntTensor<Flex>)

Gets the minimum elements and corresponding indices along a dimension.

fn int_any(tensor: IntTensor<Flex>, out_dtype: BoolDType) -> BoolTensor<Flex>

Tests if any element in the int tensor evaluates to True.

fn int_any_dim( tensor: IntTensor<Flex>, dim: usize, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Tests if any element in the int tensor evaluates to True along a given dimension dim.

fn int_all(tensor: IntTensor<Flex>, out_dtype: BoolDType) -> BoolTensor<Flex>

Tests if all elements in the int tensor evaluate to True.

fn int_all_dim( tensor: IntTensor<Flex>, dim: usize, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Tests if all elements in the int tensor evaluate to True along a given dimension dim.

fn int_powi(lhs: IntTensor<Flex>, rhs: IntTensor<Flex>) -> IntTensor<Flex>

Element-wise power with a IntTensor.

fn int_zeros( shape: Shape, _device: &Device<Flex>, dtype: IntDType, ) -> IntTensor<Flex>

Creates a tensor of zeros.

fn int_ones( shape: Shape, _device: &Device<Flex>, dtype: IntDType, ) -> IntTensor<Flex>

Creates a tensor of ones.

fn int_full( shape: Shape, fill_value: Scalar, _device: &Device<Flex>, dtype: IntDType, ) -> IntTensor<Flex>

Creates a tensor filled with given value.

fn int_transpose(tensor: IntTensor<Flex>) -> IntTensor<Flex>

Transposes an int tensor.

fn int_repeat_dim( tensor: IntTensor<Flex>, dim: usize, times: usize, ) -> IntTensor<Flex>

Repeats the tensor along the given dimension the given number of times.

fn int_not_equal( lhs: IntTensor<Flex>, rhs: IntTensor<Flex>, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise non-equality comparison.

fn int_not_equal_elem( lhs: IntTensor<Flex>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Flex>

Element-wise non-equality comparison with a scalar.

fn int_sort( tensor: IntTensor<Flex>, dim: usize, descending: bool, ) -> IntTensor<Flex>

Sort the elements of the input tensor by value along a given dimension.

fn int_sort_with_indices( tensor: IntTensor<Flex>, dim: usize, descending: bool, ) -> (IntTensor<Flex>, IntTensor<Flex>)

Sort the elements of the input tensor by value along a given dimension.

fn int_argsort( tensor: IntTensor<Flex>, dim: usize, descending: bool, ) -> IntTensor<Flex>

Returns the indices that sort the elements of the input tensor by value along a given dimension.

fn int_powi_scalar(lhs: IntTensor<Flex>, rhs: Scalar) -> IntTensor<Flex>

Element-wise power with a scalar.

fn int_powi_scalar_impl(lhs: IntTensor<Flex>, rhs: Scalar) -> IntTensor<Flex>

Element-wise power with a scalar.

fn int_max_abs(tensor: IntTensor<Flex>) -> IntTensor<Flex>

Gets the maximum absolute element in the tensor.

fn int_max_abs_dim(tensor: IntTensor<Flex>, dim: usize) -> IntTensor<Flex>

Gets the maximum absolute element in the tensor along a dimension.

fn int_arange( range: Range<i64>, _device: &Device<Flex>, dtype: IntDType, ) -> IntTensor<Flex>

Creates a new tensor with values from the given range.

fn int_arange_step( range: Range<i64>, step: usize, _device: &Device<Flex>, dtype: IntDType, ) -> IntTensor<Flex>

Creates a new tensor with values from the given range with the given step size.

fn int_mask_select( tensor: <B as BackendTypes>::IntTensorPrimitive, mask: <B as BackendTypes>::BoolTensorPrimitive, ) -> impl Future<Output = <B as BackendTypes>::IntTensorPrimitive> + Send + 'static

Selects the elements of the tensor where the mask is true, returned as a 1D tensor.

fn int_scatter( dim: usize, tensor: <B as BackendTypes>::IntTensorPrimitive, indices: <B as BackendTypes>::IntTensorPrimitive, value: <B as BackendTypes>::IntTensorPrimitive, update: IndexingUpdateOp, ) -> <B as BackendTypes>::IntTensorPrimitive

Scatter elements into a tensor using the specified update operation.

fn int_select_assign( tensor: <B as BackendTypes>::IntTensorPrimitive, dim: usize, indices: <B as BackendTypes>::IntTensorPrimitive, value: <B as BackendTypes>::IntTensorPrimitive, update: IndexingUpdateOp, ) -> <B as BackendTypes>::IntTensorPrimitive

Assign selected elements along a dimension using the specified update operation.

fn int_square( tensor: <B as BackendTypes>::IntTensorPrimitive, ) -> <B as BackendTypes>::IntTensorPrimitive

Element-wise square with a IntTensor.

fn int_argtopk( tensor: <B as BackendTypes>::IntTensorPrimitive, dim: usize, k: usize, ) -> <B as BackendTypes>::IntTensorPrimitive

Gets the indices of the k maximum elements along a dimension. If two elements share the same value, it will be ordered by the lowest coordinate

fn int_topk( tensor: <B as BackendTypes>::IntTensorPrimitive, dim: usize, k: usize, ) -> <B as BackendTypes>::IntTensorPrimitive

Gets the values of the k maximum elements along a dimension.

fn int_topk_with_indices( tensor: <B as BackendTypes>::IntTensorPrimitive, dim: usize, k: usize, ) -> (<B as BackendTypes>::IntTensorPrimitive, <B as BackendTypes>::IntTensorPrimitive)

Gets the values of the k maximum elements along a dimension, and their indices.

impl ModuleOps<Flex> for Flex

fn conv1d( x: FloatTensor<Flex>, weight: FloatTensor<Flex>, bias: Option<FloatTensor<Flex>>, options: ConvOptions<1>, ) -> FloatTensor<Flex>

One dimensional convolution.

fn conv2d( x: FloatTensor<Flex>, weight: FloatTensor<Flex>, bias: Option<FloatTensor<Flex>>, options: ConvOptions<2>, ) -> FloatTensor<Flex>

Two dimensional convolution.

fn deform_conv2d( x: FloatTensor<Flex>, offset: FloatTensor<Flex>, weight: FloatTensor<Flex>, mask: Option<FloatTensor<Flex>>, bias: Option<FloatTensor<Flex>>, options: DeformConvOptions<2>, ) -> FloatTensor<Flex>

Two dimensional deformable convolution.

fn deform_conv2d_backward( x: FloatTensor<Flex>, offset: FloatTensor<Flex>, weight: FloatTensor<Flex>, mask: Option<FloatTensor<Flex>>, bias: Option<FloatTensor<Flex>>, output_grad: FloatTensor<Flex>, options: DeformConvOptions<2>, ) -> DeformConv2dBackward<Flex>

Backward pass for the deform_conv2d operation.

fn conv3d( x: FloatTensor<Flex>, weight: FloatTensor<Flex>, bias: Option<FloatTensor<Flex>>, options: ConvOptions<3>, ) -> FloatTensor<Flex>

Three dimensional convolution.

fn conv_transpose1d( x: FloatTensor<Flex>, weight: FloatTensor<Flex>, bias: Option<FloatTensor<Flex>>, options: ConvTransposeOptions<1>, ) -> FloatTensor<Flex>

One dimensional transposed convolution.

fn conv_transpose2d( x: FloatTensor<Flex>, weight: FloatTensor<Flex>, bias: Option<FloatTensor<Flex>>, options: ConvTransposeOptions<2>, ) -> FloatTensor<Flex>

Two dimensional transposed convolution.

fn conv_transpose3d( x: FloatTensor<Flex>, weight: FloatTensor<Flex>, bias: Option<FloatTensor<Flex>>, options: ConvTransposeOptions<3>, ) -> FloatTensor<Flex>

Three dimensional transposed convolution.

fn avg_pool2d( x: FloatTensor<Flex>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], count_include_pad: bool, ceil_mode: bool, ) -> FloatTensor<Flex>

Two dimensional avg pooling.

fn avg_pool2d_backward( x: FloatTensor<Flex>, grad: FloatTensor<Flex>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], count_include_pad: bool, _divisor_override: bool, ) -> FloatTensor<Flex>

Backward pass for the avg pooling 2d operation.

fn adaptive_avg_pool2d( x: FloatTensor<Flex>, output_size: [usize; 2], ) -> FloatTensor<Flex>

Two dimensional adaptive avg pooling.

fn adaptive_avg_pool2d_backward( x: FloatTensor<Flex>, grad: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Backward pass for the adaptive avg pooling 2d operation.

fn adaptive_avg_pool3d( x: FloatTensor<Flex>, output_size: [usize; 3], ) -> FloatTensor<Flex>

Three dimensional adaptive avg pooling.

fn adaptive_avg_pool3d_backward( x: FloatTensor<Flex>, grad: FloatTensor<Flex>, ) -> FloatTensor<Flex>

Backward pass for the adaptive avg pooling 3d operation.

fn max_pool2d( x: FloatTensor<Flex>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, ) -> FloatTensor<Flex>

Two dimensional max pooling.

fn max_pool2d_with_indices( x: FloatTensor<Flex>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, indices_dtype: IntDType, ) -> MaxPool2dWithIndices<Flex>

Two dimensional max pooling with indices.

fn max_pool2d_with_indices_backward( x: FloatTensor<Flex>, _kernel_size: [usize; 2], _stride: [usize; 2], _padding: [usize; 2], _dilation: [usize; 2], _ceil_mode: bool, output_grad: FloatTensor<Flex>, indices: IntTensor<Flex>, ) -> MaxPool2dBackward<Flex>

Backward pass for the max pooling 2d operation.

fn interpolate( x: FloatTensor<Flex>, output_size: [usize; 2], options: InterpolateOptions, ) -> FloatTensor<Flex>

Down/up samples the input.

fn interpolate_backward( x: FloatTensor<Flex>, grad: FloatTensor<Flex>, output_size: [usize; 2], options: InterpolateOptions, ) -> FloatTensor<Flex>

Backward pass for the interpolate operation.

fn attention( query: FloatTensor<Flex>, key: FloatTensor<Flex>, value: FloatTensor<Flex>, mask: Option<BoolTensor<Flex>>, attn_bias: Option<FloatTensor<Flex>>, options: AttentionModuleOptions, ) -> FloatTensor<Flex>

Computes scaled dot-product attention: softmax(QKᵗ * scale) · V, where scale defaults to 1/sqrt(head_dim). Optionally applies masking, additive bias, causal masking, and softcap to the attention scores.

fn rfft( signal: FloatTensor<Flex>, dim: usize, n: Option<usize>, ) -> (FloatTensor<Flex>, FloatTensor<Flex>)

Real-valued FFT with optional size parameter.

fn irfft( spectrum_re: FloatTensor<Flex>, spectrum_im: FloatTensor<Flex>, dim: usize, n: Option<usize>, ) -> FloatTensor<Flex>

Inverse real-valued FFT with optional output size.

fn embedding( weights: FloatTensor<Flex>, indices: IntTensor<Flex>, ) -> FloatTensor<Flex>

Embedding operation.

fn layer_norm( tensor: FloatTensor<Flex>, gamma: FloatTensor<Flex>, beta: Option<FloatTensor<Flex>>, epsilon: f64, ) -> FloatTensor<Flex>

Applies Layer Normalization over the last dimension of the input tensor.

fn embedding_backward( weights: FloatTensor<Flex>, output_grad: FloatTensor<Flex>, indices: IntTensor<Flex>, ) -> FloatTensor<Flex>

Embedding backward operation.

fn linear( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, bias: Option<<B as BackendTypes>::FloatTensorPrimitive>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Linear transformation.

fn linear_x_backward( weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for linear, returning the gradient for x.

fn linear_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for linear, returning the gradient for weight.

fn linear_bias_backward( output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for linear, returning the gradient for bias.

fn conv1d_x_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<1>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv1d operation, returning the gradient for x.

fn conv1d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<1>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv1d operation, returning the gradient for weight.

fn conv1d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv1d operation, returning the gradient for bias.

fn conv2d_x_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<2>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv2d operation, returning the gradient for x.

fn conv2d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<2>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv2d operation, returning the gradient for weight.

fn conv2d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv2d operation, returning the gradient for bias.

fn conv3d_x_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<3>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv3d operation, returning the gradient for x.

fn conv3d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<3>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv3d operation, returning the gradient for weight.

fn conv3d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv3d operation, returning the gradient for bias.

fn conv_transpose1d_x_backward( weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<1>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 1d operation, returning the gradient for x.

fn conv_transpose1d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<1>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 1d operation, returning the gradient for weight.

fn conv_transpose1d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 1d operation, returning the gradient for bias.

fn conv_transpose2d_x_backward( weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<2>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 2d operation, returning the gradient for x.

fn conv_transpose2d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<2>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 2d operation, returning the gradient for weight.

fn conv_transpose2d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 2d operation, returning the gradient for bias.

fn conv_transpose3d_x_backward( weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<3>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 3d operation, returning the gradient for x.

fn conv_transpose3d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<3>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 3d operation, returning the gradient for weight.

fn conv_transpose3d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 3d operation, returning the gradient for bias.

fn unfold4d( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: [usize; 2], options: UnfoldOptions, ) -> <B as BackendTypes>::FloatTensorPrimitive

Four-dimensional unfolding.

fn fold4d( x: <B as BackendTypes>::FloatTensorPrimitive, output_size: [usize; 2], kernel_size: [usize; 2], options: UnfoldOptions, ) -> <B as BackendTypes>::FloatTensorPrimitive

Four dimensional fold (col2im), the adjoint of unfold4d.

fn avg_pool1d( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: usize, stride: usize, padding: usize, count_include_pad: bool, ceil_mode: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

One dimensional avg pooling.

fn avg_pool1d_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: usize, stride: usize, padding: usize, count_include_pad: bool, ceil_mode: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the avg pooling 1d operation.

fn adaptive_avg_pool1d( x: <B as BackendTypes>::FloatTensorPrimitive, output_size: usize, ) -> <B as BackendTypes>::FloatTensorPrimitive

One dimensional adaptive avg pooling.

fn adaptive_avg_pool1d_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the adaptive avg pooling 1d operation.

fn max_pool1d( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: usize, stride: usize, padding: usize, dilation: usize, ceil_mode: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

One dimensional max pooling.

fn max_pool1d_with_indices( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: usize, stride: usize, padding: usize, dilation: usize, ceil_mode: bool, indices_dtype: IntDType, ) -> MaxPool1dWithIndices<B>

One dimensional max pooling with indices.

fn max_pool1d_with_indices_backward( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: usize, stride: usize, padding: usize, dilation: usize, ceil_mode: bool, output_grad: <B as BackendTypes>::FloatTensorPrimitive, indices: <B as BackendTypes>::IntTensorPrimitive, ) -> MaxPool1dBackward<B>

Backward pass for the max pooling 1d operation.

fn ctc_loss( log_probs: <B as BackendTypes>::FloatTensorPrimitive, targets: <B as BackendTypes>::IntTensorPrimitive, input_lengths: <B as BackendTypes>::IntTensorPrimitive, target_lengths: <B as BackendTypes>::IntTensorPrimitive, blank: usize, ) -> <B as BackendTypes>::FloatTensorPrimitive

Computes the Connectionist Temporal Classification (CTC) loss.

fn has_ctc_loss_backward() -> bool

Returns true if this backend implements ctc_loss_backward natively.

fn ctc_loss_backward( _log_probs: <B as BackendTypes>::FloatTensorPrimitive, _targets: <B as BackendTypes>::IntTensorPrimitive, _input_lengths: <B as BackendTypes>::IntTensorPrimitive, _target_lengths: <B as BackendTypes>::IntTensorPrimitive, _grad_loss: <B as BackendTypes>::FloatTensorPrimitive, _blank: usize, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for ctc_loss: gradient w.r.t. log_probs.

impl QTensorOps<Flex> for Flex

fn q_from_data( data: TensorData, _device: &Device<Flex>, ) -> QuantizedTensor<Flex>

Creates a new tensor from the data structure.

fn quantize_dynamic( tensor: FloatTensor<Flex>, scheme: &QuantScheme, ) -> QuantizedTensor<Flex>

Dynamically convert the tensor to a lower precision data type based on the quantization scheme.

fn quantize( tensor: FloatTensor<Flex>, scheme: &QuantScheme, qparams: QuantizationParametersPrimitive<Flex>, ) -> QuantizedTensor<Flex>

Convert the tensor to a lower precision data type based on the quantization scheme and parameters.

fn dequantize( tensor: QuantizedTensor<Flex>, dtype: FloatDType, ) -> FloatTensor<Flex>

Convert the tensor back to a higher precision data type.

fn q_to_device( tensor: QuantizedTensor<Flex>, _device: &Device<Flex>, ) -> QuantizedTensor<Flex>

Moves the tensor to the given device.

fn q_reshape( tensor: QuantizedTensor<Flex>, shape: Shape, ) -> QuantizedTensor<Flex>

Reshapes a tensor.

async fn q_into_data( tensor: QuantizedTensor<Flex>, ) -> Result<TensorData, ExecutionError>

Converts the tensor to a data structure.

fn q_swap_dims( tensor: QuantizedTensor<Flex>, dim1: usize, dim2: usize, ) -> QuantizedTensor<Flex>

Swaps two dimensions of a tensor.

fn q_permute( tensor: QuantizedTensor<Flex>, axes: &[usize], ) -> QuantizedTensor<Flex>

Permutes the dimensions of a tensor.

fn q_flip( tensor: QuantizedTensor<Flex>, axes: &[usize], ) -> QuantizedTensor<Flex>

Reverse the order of elements in a tensor along the given axes.

fn q_expand( tensor: QuantizedTensor<Flex>, shape: Shape, ) -> QuantizedTensor<Flex>

Broadcasts the tensor to the given shape.

fn q_select( tensor: QuantizedTensor<Flex>, dim: usize, indices: IntTensor<Flex>, ) -> QuantizedTensor<Flex>

Select tensor elements along the given dimension corresponding for the given indices.

fn q_slice( tensor: QuantizedTensor<Flex>, slices: &[Slice], ) -> QuantizedTensor<Flex>

Select tensor elements corresponding to the given slices.

fn q_argmax( tensor: QuantizedTensor<Flex>, dim: usize, out_dtype: IntDType, ) -> IntTensor<Flex>

Gets the indices of the maximum elements of a tensor along an axis.

fn q_argmin( tensor: QuantizedTensor<Flex>, dim: usize, out_dtype: IntDType, ) -> IntTensor<Flex>

Gets the indices of the minimum elements of a tensor along an axis.

fn q_gather( dim: usize, tensor: QuantizedTensor<Flex>, indices: IntTensor<Flex>, ) -> QuantizedTensor<Flex>

Gather elements from a tensor.

fn q_detach( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Detaches a tensor from the computation graph.

fn q_set_require_grad( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, _require_grad: bool, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Sets the require_grad flag of a tensor.

fn q_is_require_grad( _tensor: &<B as BackendTypes>::QuantizedTensorPrimitive, ) -> bool

Returns the require_grad flag of a tensor.

fn q_transpose( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Transposes a tensor.

fn q_repeat_dim( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, times: usize, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Repeat the tensor along the given dimension.

fn q_add( lhs: <B as BackendTypes>::QuantizedTensorPrimitive, rhs: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Adds two tensors together.

fn q_add_scalar( lhs: <B as BackendTypes>::QuantizedTensorPrimitive, rhs: Scalar, ) -> TensorPrimitive<B>

Adds a scalar to a tensor.

fn q_clamp_min( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, min: Scalar, ) -> TensorPrimitive<B>

Clamps a tensor under a minimum value.

fn q_clamp_max( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, max: Scalar, ) -> TensorPrimitive<B>

Clamps a tensor over a maximum value.

fn q_clamp( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, min: Scalar, max: Scalar, ) -> TensorPrimitive<B>

Clamps a tensor between a minimum and maximum value.

fn q_sub( lhs: <B as BackendTypes>::QuantizedTensorPrimitive, rhs: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Subtracts two tensors.

fn q_sub_scalar( lhs: <B as BackendTypes>::QuantizedTensorPrimitive, rhs: Scalar, ) -> TensorPrimitive<B>

Subtracts a scalar from a tensor.

fn q_mul( lhs: <B as BackendTypes>::QuantizedTensorPrimitive, rhs: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Multiplies two tensors together element-wise.

fn q_mul_scalar( lhs: <B as BackendTypes>::QuantizedTensorPrimitive, rhs: Scalar, ) -> TensorPrimitive<B>

Multiplies a tensor by a scalar.

fn q_div( lhs: <B as BackendTypes>::QuantizedTensorPrimitive, rhs: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Divides two tensors element-wise.

fn q_div_scalar( lhs: <B as BackendTypes>::QuantizedTensorPrimitive, rhs: Scalar, ) -> TensorPrimitive<B>

Divides a tensor by a scalar.

fn q_matmul( lhs: TensorPrimitive<B>, rhs: TensorPrimitive<B>, ) -> TensorPrimitive<B>

Multiplies two tensors together using matrix multiplication.

fn q_neg( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Negates a tensor element-wise.

fn q_recip( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Calculates the reciprocals element-wise

fn q_sum( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Sum of all elements in a tensor.

fn q_sum_dim( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, ) -> TensorPrimitive<B>

Sum of all elements in a tensor along a dimension.

fn q_prod( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Product of all elements in a tensor.

fn q_prod_dim( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, ) -> TensorPrimitive<B>

Product of all elements in a tensor along a dimension.

fn q_mean( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Mean of all elements in a tensor.

fn q_mean_dim( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, ) -> TensorPrimitive<B>

Mean of all elements in a tensor along a dimension.

fn q_cumsum( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, ) -> TensorPrimitive<B>

Computes the cumulative sum of elements along a dimension.

fn q_cumprod( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, ) -> TensorPrimitive<B>

Computes the cumulative product of elements along a dimension.

fn q_cummin( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, ) -> TensorPrimitive<B>

Computes the cumulative minimum of elements along a dimension.

fn q_cummax( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, ) -> TensorPrimitive<B>

Computes the cumulative maximum of elements along a dimension.

fn q_exp( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Returns a new tensor with exponential values.

fn q_log( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Returns a new tensor with natural logarithm values.

fn q_log1p( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Returns a new tensor with logarithm values of (1 + Xi).

fn q_powf( lhs: <B as BackendTypes>::QuantizedTensorPrimitive, rhs: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Element-wise power with another tensor.

fn q_powi( lhs: <B as BackendTypes>::QuantizedTensorPrimitive, rhs: <B as BackendTypes>::IntTensorPrimitive, ) -> TensorPrimitive<B>

Element-wise power with an IntTensor.

fn q_powi_scalar( lhs: <B as BackendTypes>::QuantizedTensorPrimitive, rhs: Scalar, ) -> TensorPrimitive<B>

Element-wise power with an int scalar.

fn q_powf_scalar( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, value: Scalar, ) -> TensorPrimitive<B>

Element-wise power with a float scalar.

fn q_sqrt( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Returns a new tensor with square root values.

fn q_abs( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Returns a new tensor with absolute values.

fn q_cos( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Returns a new tensor with cosine values.

fn q_sin( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Returns a new tensor with sine values.

fn q_tan( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Returns a new tensor with tangent values.

fn q_cosh( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Returns a new tensor with hyperbolic cosine values.

fn q_sinh( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Returns a new tensor with hyperbolic sine values.

fn q_tanh( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Returns a new tensor with hyperbolic tangent values.

fn q_erf( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> TensorPrimitive<B>

Returns a new tensor with the error function values.

fn q_cat( tensors: Vec<<B as BackendTypes>::QuantizedTensorPrimitive>, dim: usize, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Concatenates tensors along a dimension.

fn q_argtopk( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, k: usize, out_dtype: IntDType, ) -> <B as BackendTypes>::IntTensorPrimitive

Gets the indices of the k maximum elements of a tensor along an axis. If two elements are equals, order them by the lowest indices

fn q_topk( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, k: usize, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Gets the values of the k maximum elements of a tensor along an axis.

fn q_topk_with_indices( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, k: usize, out_dtype: IntDType, ) -> (<B as BackendTypes>::QuantizedTensorPrimitive, <B as BackendTypes>::IntTensorPrimitive)

Gets the values of the k maximum elements of a tensor along an axis, and their indices.

fn q_max( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Gets the maximum element of a tensor.

fn q_max_dim( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Gets the maximum elements of a tensor along an axis.

fn q_max_dim_with_indices( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, out_dtype: IntDType, ) -> (<B as BackendTypes>::QuantizedTensorPrimitive, <B as BackendTypes>::IntTensorPrimitive)

Gets the maximum elements of a tensor along an axis and their indices.

fn q_min( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Gets the minimum element of a tensor.

fn q_min_dim( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Gets the minimum elements of a tensor along an axis.

fn q_min_dim_with_indices( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, out_dtype: IntDType, ) -> (<B as BackendTypes>::QuantizedTensorPrimitive, <B as BackendTypes>::IntTensorPrimitive)

Gets the minimum elements of a tensor along an axis and their indices.

fn q_max_abs( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Gets the maximum element of a tensor.

fn q_max_abs_dim( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Gets the maximum elements of a tensor along an axis.

fn q_any( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, out_dtype: BoolStore, ) -> <B as BackendTypes>::BoolTensorPrimitive

Tests if any element in the tensor evaluates to True.

fn q_any_dim( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, out_dtype: BoolStore, ) -> <B as BackendTypes>::BoolTensorPrimitive

Tests if any element in the float tensor evaluates to True along a given dimension dim.

fn q_all( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, out_dtype: BoolStore, ) -> <B as BackendTypes>::BoolTensorPrimitive

Tests if all elements in the tensor evaluate to True.

fn q_all_dim( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, out_dtype: BoolStore, ) -> <B as BackendTypes>::BoolTensorPrimitive

Tests if all elements in the tensor evaluate to True along a given dimension dim.

fn q_sort( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, descending: bool, ) -> <B as BackendTypes>::QuantizedTensorPrimitive

Sort the elements of the input tensor by value in along a given dimension.

fn q_sort_with_indices( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, descending: bool, out_dtype: IntDType, ) -> (<B as BackendTypes>::QuantizedTensorPrimitive, <B as BackendTypes>::IntTensorPrimitive)

Sort the elements of the input tensor by value in along a given dimension.

fn q_argsort( tensor: <B as BackendTypes>::QuantizedTensorPrimitive, dim: usize, descending: bool, out_dtype: IntDType, ) -> <B as BackendTypes>::IntTensorPrimitive

Returns the indices that sort the elements of the input tensor by value along a given dimension.

impl TransactionOps<Flex> for Flex

fn tr_execute( transaction: TransactionPrimitive<B>, ) -> impl Future<Output = Result<TransactionPrimitiveData, ExecutionError>> + Send

Executes a transaction and return its data.