Struct Flex 

pub struct Flex {}

Available on default_backend or crate feature flex only.

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: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Applies the ReLU activation function.

fn relu_backward( output: <Flex as BackendTypes>::FloatTensorPrimitive, grad: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Applies the ReLU activation function backward.

fn leaky_relu( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, negative_slope: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Applies the LeakyReLU activation function.

fn prelu( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, alpha: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Applies the PReLu activation function.

fn gelu( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Applies the Gelu activation function.

fn gelu_backward( x: <Flex as BackendTypes>::FloatTensorPrimitive, grad: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Applies the Gelu activation function backward.

fn sigmoid( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Applies the Sigmoid activation function.

fn sigmoid_backward( output: <Flex as BackendTypes>::FloatTensorPrimitive, grad: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Applies the Sigmoid activation function backward.

fn hard_sigmoid( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, alpha: Scalar, beta: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Applies the hard Sigmoid activation function.

fn log_sigmoid( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Applies the LogSigmoid activation function.

fn log_sigmoid_backward( x: <Flex as BackendTypes>::FloatTensorPrimitive, grad: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Applies the LogSigmoid activation function backward.

fn softmax( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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: &<Flex as BackendTypes>::Device) -> String

Name of the backend.

fn seed(_device: &<Flex as BackendTypes>::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: &<Flex as BackendTypes>::Device, dtype: DType, ) -> EnumSet<DTypeUsage>

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

fn flush(_device: &<Flex as BackendTypes>::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<<Flex as BackendIr>::Handle>) -> FlexTensor

Convert a handle to a float tensor.

fn int_tensor(handle: TensorHandle<<Flex as BackendIr>::Handle>) -> FlexTensor

Convert a handle to an int tensor.

fn bool_tensor(handle: TensorHandle<<Flex as BackendIr>::Handle>) -> FlexTensor

Convert a handle to a bool tensor.

fn quantized_tensor( handle: TensorHandle<<Flex as BackendIr>::Handle>, ) -> FlexQTensor

Convert a handle to a quantized tensor.

fn float_tensor_handle(tensor: FlexTensor) -> <Flex as BackendIr>::Handle

Convert a float tensor to a handle.

fn int_tensor_handle(tensor: FlexTensor) -> <Flex as BackendIr>::Handle

Convert an int tensor to a handle.

fn bool_tensor_handle(tensor: FlexTensor) -> <Flex as BackendIr>::Handle

Convert a bool tensor to a handle.

fn quantized_tensor_handle(tensor: FlexQTensor) -> <Flex as BackendIr>::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: &<Flex as BackendTypes>::Device, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Creates a tensor from the data structure.

async fn bool_into_data( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> Result<TensorData, ExecutionError>

Converts the tensor to a data structure.

fn bool_to_device( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, _device: &<Flex as BackendTypes>::Device, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Moves the tensor to the device.

fn bool_cat( tensors: Vec<<Flex as BackendTypes>::BoolTensorPrimitive>, dim: usize, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Concatenates the tensors along the given dimension.

fn bool_reshape( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, shape: Shape, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Reshapes the tensor.

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

Gets the values from the tensor for the given ranges.

fn bool_empty( shape: Shape, _device: &<Flex as BackendTypes>::Device, dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Creates a new bool tensor.

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

Sets the values in the tensor for the given ranges.

fn bool_into_int( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, out_dtype: IntDType, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Converts bool tensor to int tensor.

fn bool_into_float( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, out_dtype: FloatDType, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Converts bool tensor to float tensor.

fn bool_swap_dims( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, dim1: usize, dim2: usize, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Swaps two dimensions of a bool tensor.

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

Permutes the dimensions of a tensor.

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

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

fn bool_equal( lhs: <Flex as BackendTypes>::BoolTensorPrimitive, rhs: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Equates the two tensors.

fn bool_not( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Inverses boolean values.

fn bool_and( lhs: <Flex as BackendTypes>::BoolTensorPrimitive, rhs: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn bool_or( lhs: <Flex as BackendTypes>::BoolTensorPrimitive, rhs: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn bool_xor( lhs: <Flex as BackendTypes>::BoolTensorPrimitive, rhs: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise exclusive or.

fn bool_expand( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, shape: Shape, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Broadcasts the bool tensor to the given shape.

fn bool_zeros( shape: Shape, device: &<Flex as BackendTypes>::Device, dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Creates a new bool tensor filled false.

fn bool_ones( shape: Shape, _device: &<Flex as BackendTypes>::Device, dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Creates a new bool tensor filled true.

fn bool_mask_where( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, mask: <Flex as BackendTypes>::BoolTensorPrimitive, value: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn bool_mask_fill( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, mask: <Flex as BackendTypes>::BoolTensorPrimitive, value: Scalar, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn bool_gather( dim: usize, tensor: <Flex as BackendTypes>::BoolTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Gather elements from the tensor at the given indices.

fn bool_scatter_or( dim: usize, tensor: <Flex as BackendTypes>::BoolTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, value: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn bool_equal_elem( lhs: <Flex as BackendTypes>::BoolTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise equality comparison with a scalar.

fn bool_unfold( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, dim: usize, size: usize, step: usize, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Unfold windows along a dimension.

fn bool_not_equal( lhs: <Flex as BackendTypes>::BoolTensorPrimitive, rhs: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise non-equality comparison.

fn bool_not_equal_elem( lhs: <Flex as BackendTypes>::BoolTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise non-equality comparison with a scalar.

fn bool_any( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn bool_any_dim( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn bool_all( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn bool_all_dim( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn bool_select( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, dim: usize, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn bool_select_or( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, dim: usize, indices: <Flex as BackendTypes>::IntTensorPrimitive, value: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn bool_transpose( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Transposes a bool tensor.

fn bool_repeat_dim( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, dim: usize, times: usize, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

async fn bool_argwhere( tensor: <Flex as BackendTypes>::BoolTensorPrimitive, out_dtype: IntDType, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

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<(), Error>

Formats the value using the given formatter.

impl Default for Flex

fn default() -> Flex

Returns the “default value” for a type.

impl DispatchKindConversion<Flex> for DispatchTensor

fn try_into_backend( tensor: DispatchTensor, ) -> Result<BackendTensor<Flex>, String>

Attempts to extract a backend-specific BackendTensor wrapper from a generic, dynamically-routed DispatchTensor.

fn from_backend(tensor: BackendTensor<Flex>) -> DispatchTensor

Encapsulates a backend-specific tensor variant back into a globally routing DispatchTensor.

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: &<Flex as BackendTypes>::Device, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Creates a new tensor from the data structure.

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

Creates a new tensor with random values.

async fn float_into_data( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> Result<TensorData, ExecutionError>

Converts the tensor to a data structure.

fn float_to_device( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, _device: &<Flex as BackendTypes>::Device, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Moves the tensor to the given device.

fn float_detach( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Detaches a tensor from the computation graph.

fn float_into_int( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, out_dtype: IntDType, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Converts float tensor to int tensor.

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

Creates an empty tensor with the given shape.

fn float_add( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Adds two tensors together.

fn float_add_scalar( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Adds a scalar to a tensor.

fn float_sub( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Subtracts two tensors.

fn float_sub_scalar( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Subtracts a scalar from a tensor.

fn float_mul( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Multiplies two tensors together element-wise.

fn float_mul_scalar( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Multiplies a tensor by a scalar.

fn float_div( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Divides two tensors element-wise.

fn float_div_scalar( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Divides a tensor by a scalar.

fn float_remainder( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_remainder_scalar( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Computes the modulus of a tensor given a scalar.

fn float_matmul( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Multiplies two tensors together using matrix multiplication.

fn float_cross( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_recip( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Calculates the reciprocals element-wise

fn float_swap_dims( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim1: usize, dim2: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Swaps two dimensions of a tensor.

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

Permutes the dimensions of a tensor.

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

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

fn float_cat( tensors: Vec<<Flex as BackendTypes>::FloatTensorPrimitive>, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Concatenates tensors along a dimension.

fn float_reshape( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, shape: Shape, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Reshapes a tensor.

fn float_gather( dim: usize, tensor: <Flex as BackendTypes>::FloatTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Gather elements from a tensor.

fn float_scatter_add( dim: usize, tensor: <Flex as BackendTypes>::FloatTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, value: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Scatter elements into a tensor using sum reduction.

fn float_scatter_nd( data: <Flex as BackendTypes>::FloatTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, values: <Flex as BackendTypes>::FloatTensorPrimitive, reduction: IndexingUpdateOp, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_gather_nd( data: <Flex as BackendTypes>::FloatTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_select( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_select_add( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, indices: <Flex as BackendTypes>::IntTensorPrimitive, value: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

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

Select tensor elements corresponding to the given slices.

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

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

fn float_mask_where( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, mask: <Flex as BackendTypes>::BoolTensorPrimitive, value: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_mask_fill( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, mask: <Flex as BackendTypes>::BoolTensorPrimitive, value: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_equal( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Equal comparison of two tensors.

fn float_equal_elem( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Equal comparison of a tensor and a scalar.

fn float_greater( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Greater than comparison of two tensors.

fn float_greater_elem( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Greater than comparison of a tensor and a scalar.

fn float_greater_equal( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Greater than or equal comparison of two tensors.

fn float_greater_equal_elem( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn float_lower( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Less than comparison of two tensors.

fn float_lower_elem( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Less than comparison of a tensor and a scalar.

fn float_lower_equal( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Less than or equal comparison of two tensors.

fn float_lower_equal_elem( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn float_not_equal( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise non-equality comparison.

fn float_not_equal_elem( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise non-equality comparison with a scalar.

fn float_neg( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Negates a tensor element-wise.

fn float_clamp( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, min: Scalar, max: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Clamps a tensor between a minimum and maximum value.

fn float_clamp_min( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, min: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Clamps a tensor under a minimum value.

fn float_clamp_max( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, max: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Clamps a tensor over a maximum value.

fn float_sign( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns the signs of the float tensor.

fn float_mean( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Mean of all elements in a tensor.

fn float_max( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Gets the maximum element of a tensor.

fn float_max_dim( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Gets the maximum elements of a tensor along an axis.

fn float_min( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Gets the minimum element of a tensor.

fn float_min_dim( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Gets the minimum elements of a tensor along an axis.

fn float_max_dim_with_indices( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, indices_dtype: IntDType, ) -> (<Flex as BackendTypes>::FloatTensorPrimitive, <Flex as BackendTypes>::IntTensorPrimitive)

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

fn float_min_dim_with_indices( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, indices_dtype: IntDType, ) -> (<Flex as BackendTypes>::FloatTensorPrimitive, <Flex as BackendTypes>::IntTensorPrimitive)

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

fn float_any( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn float_any_dim( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn float_all( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn float_all_dim( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn float_sum( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Sum of all elements in a tensor.

fn float_sum_dim( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Sum of all elements in a tensor along a dimension.

fn float_mean_dim( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Mean of all elements in a tensor along a dimension.

fn float_prod( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Product of all elements in a tensor.

fn float_prod_dim( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Product of all elements in a tensor along a dimension.

fn float_cumsum( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Computes the cumulative sum of elements along a dimension.

fn float_cumprod( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Computes the cumulative product of elements along a dimension.

fn float_cummin( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Computes the cumulative minimum of elements along a dimension.

fn float_cummax( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Computes the cumulative maximum of elements along a dimension.

fn float_cast( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dtype: FloatDType, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Converts a tensor to another floating point data type.

fn float_exp( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with exponential values.

fn float_log( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with natural logarithm values.

fn float_log1p( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_powf( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Element-wise power with a FloatTensor.

fn float_powf_scalar_impl( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, value: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_sqrt( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with square root values.

fn float_abs( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with absolute values.

fn float_cos( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with cosine values.

fn float_sin( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with sine values.

fn float_tan( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with tangent values.

fn float_cosh( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with hyperbolic cosine values.

fn float_sinh( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with hyperbolic sine values.

fn float_tanh( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with hyperbolic tangent values.

fn float_acos( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with inverse cosine values.

fn float_acosh( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with inverse hyperbolic cosine values.

fn float_asin( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with inverse sine values.

fn float_asinh( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with inverse hyperbolic sine values.

fn float_atan( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with the inverse tangent values.

fn float_atanh( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with the inverse hyperbolic tangent values.

fn float_atan2( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_round( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with rounded values.

fn float_floor( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with floored values.

fn float_ceil( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with ceiled values.

fn float_trunc( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with truncated values.

fn float_erf( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with the error function values.

fn float_argmax( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, out_dtype: IntDType, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

fn float_argmin( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, out_dtype: IntDType, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

fn float_expand( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, shape: Shape, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Broadcasts the float tensor to the given shape.

fn float_unfold( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, size: usize, step: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Unfold windows along a dimension.

fn float_grid_sample_2d( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, grid: <Flex as BackendTypes>::FloatTensorPrimitive, options: GridSampleOptions, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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: &<Flex as BackendTypes>::Device, dtype: FloatDType, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Creates a new tensor with zeros.

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

Creates a new tensor with ones.

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

Creates a tensor filled with given value.

fn float_transpose( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Transposes a tensor.

fn float_repeat_dim( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, times: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Repeat the tensor along the given dimension.

fn float_sort( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, descending: bool, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_sort_with_indices( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, descending: bool, indices_dtype: IntDType, ) -> (<Flex as BackendTypes>::FloatTensorPrimitive, <Flex as BackendTypes>::IntTensorPrimitive)

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

fn float_argsort( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, descending: bool, out_dtype: IntDType, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

fn float_powi( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Element-wise power with an IntTensor.

fn float_powi_scalar( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Raises a tensor to the power of an int scalar.

fn float_powf_scalar( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, value: Scalar, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_max_abs( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Gets the maximum absolute element of a tensor.

fn float_max_abs_dim( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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

fn float_is_nan( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn float_is_inf( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn float_hypot( lhs: <Flex as BackendTypes>::FloatTensorPrimitive, rhs: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Returns a new tensor with the Euclidean distance values.

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: <Flex as BackendTypes>::IntTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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: <Flex as BackendTypes>::IntTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, value: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, indices: <Flex as BackendTypes>::IntTensorPrimitive, value: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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: &<Flex as BackendTypes>::Device, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Creates a tensor from the data structure.

async fn int_into_data( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> Result<TensorData, ExecutionError>

Converts the tensor to a data structure.

fn int_to_device( tensor: <Flex as BackendTypes>::IntTensorPrimitive, _device: &<Flex as BackendTypes>::Device, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Moves the tensor to the given device.

fn int_cat( tensors: Vec<<Flex as BackendTypes>::IntTensorPrimitive>, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Concatenates the given tensors along the given dimension.

fn int_reshape( tensor: <Flex as BackendTypes>::IntTensorPrimitive, shape: Shape, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Reshapes the tensor.

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

Gets the element at the given indices.

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

Creates a new int tensor.

fn int_mask_where( tensor: <Flex as BackendTypes>::IntTensorPrimitive, mask: <Flex as BackendTypes>::BoolTensorPrimitive, value: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

fn int_mask_fill( tensor: <Flex as BackendTypes>::IntTensorPrimitive, mask: <Flex as BackendTypes>::BoolTensorPrimitive, value: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

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

Sets the values in the tensor for the given ranges.

fn int_scatter_nd( data: <Flex as BackendTypes>::IntTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, values: <Flex as BackendTypes>::IntTensorPrimitive, reduction: IndexingUpdateOp, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Multi-dimensional scatter for int tensors.

fn int_gather_nd( data: <Flex as BackendTypes>::IntTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Multi-dimensional gather for int tensors.

fn int_equal( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise equality comparison.

fn int_equal_elem( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise equality comparison with a scalar.

fn int_greater( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise greater than comparison.

fn int_greater_elem( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise greater than comparison with a scalar.

fn int_greater_equal( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise greater than or equal comparison.

fn int_greater_equal_elem( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn int_lower( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise less than comparison.

fn int_lower_elem( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise less than comparison with a scalar.

fn int_lower_equal( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise less than or equal comparison.

fn int_lower_equal_elem( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn int_add( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise addition.

fn int_add_scalar( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise addition with a scalar.

fn int_sub( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise subtraction.

fn int_sub_scalar( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise subtraction with a scalar.

fn int_mul( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise multiplication.

fn int_mul_scalar( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise multiplication with a scalar.

fn int_div( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise division.

fn int_div_scalar( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise division with a scalar.

fn int_remainder( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise modulus.

fn int_remainder_scalar( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise modulus with a scalar.

fn int_into_float( tensor: <Flex as BackendTypes>::IntTensorPrimitive, out_dtype: FloatDType, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Converts int tensor to float tensor.

fn int_swap_dims( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim1: usize, dim2: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Swaps two dimensions of an int tensor.

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

Permutes the dimensions of a tensor.

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

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

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

Creates a new int tensor with random values.

fn int_expand( tensor: <Flex as BackendTypes>::IntTensorPrimitive, shape: Shape, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Broadcasts the int tensor to the given shape.

fn int_matmul( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Multiplies two tensors together using matrix multiplication.

fn int_sum( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Sums all elements in the tensor.

fn int_sum_dim( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Sums all elements in the tensor along a dimension.

fn int_prod( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Computes the product of all elements in the tensor.

fn int_prod_dim( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

fn int_mean_dim( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

fn int_cumsum( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Computes the cumulative sum of elements along a dimension.

fn int_cumprod( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Computes the cumulative product of elements along a dimension.

fn int_cummin( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Computes the cumulative minimum of elements along a dimension.

fn int_cummax( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Computes the cumulative maximum of elements along a dimension.

fn int_argmax( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Gets the indices of the maximum elements along a dimension.

fn int_argmin( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Gets the indices of the minimum elements along a dimension.

fn int_abs( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Returns a new tensor with absolute values.

fn bitwise_and( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Bitwise AND operation for Int Tensors

fn bitwise_and_scalar( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Bitwise AND operation for Int Tensors with a scalar

fn bitwise_or( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Bitwise OR operation for Int Tensors

fn bitwise_or_scalar( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Bitwise OR operation for Int Tensors with a scalar

fn bitwise_xor( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Bitwise XOR operation for Int Tensors

fn bitwise_xor_scalar( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Bitwise XOR operation for Int Tensors with a scalar

fn bitwise_not( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Bitwise NOT operation for Int Tensors

fn bitwise_left_shift( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Bitwise left shift operation for Int Tensors

fn bitwise_left_shift_scalar( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Bitwise left shift operation for Int Tensors with a scalar

fn bitwise_right_shift( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Bitwise right shift operation for Int Tensors

fn bitwise_right_shift_scalar( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Bitwise right shift operation for Int Tensors with a scalar

fn int_cast( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dtype: IntDType, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Converts a tensor to another integer data type.

fn int_unfold( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, size: usize, step: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Unfold windows along a dimension.

fn int_neg( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise negation.

fn int_clamp( tensor: <Flex as BackendTypes>::IntTensorPrimitive, min: Scalar, max: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Clamps a tensor between a minimum and maximum value.

fn int_clamp_min( tensor: <Flex as BackendTypes>::IntTensorPrimitive, min: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Clamps a tensor under a minimum value.

fn int_clamp_max( tensor: <Flex as BackendTypes>::IntTensorPrimitive, max: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Clamps a tensor over a maximum value.

fn int_sign( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Returns the signs of the int tensor.

fn int_mean( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Computes the mean of all elements in the tensor.

fn int_max( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Gets the maximum element in the tensor.

fn int_max_dim( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Gets the maximum element in the tensor along a dimension.

fn int_min( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Gets the minimum element in the tensor.

fn int_min_dim( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Gets the minimum elements in the tensor along a dimension.

fn int_max_dim_with_indices( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> (<Flex as BackendTypes>::IntTensorPrimitive, <Flex as BackendTypes>::IntTensorPrimitive)

Gets the maximum elements and corresponding indices along a dimension.

fn int_min_dim_with_indices( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> (<Flex as BackendTypes>::IntTensorPrimitive, <Flex as BackendTypes>::IntTensorPrimitive)

Gets the minimum elements and corresponding indices along a dimension.

fn int_any( tensor: <Flex as BackendTypes>::IntTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn int_any_dim( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn int_all( tensor: <Flex as BackendTypes>::IntTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn int_all_dim( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

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

fn int_powi( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise power with a IntTensor.

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

Creates a tensor of zeros.

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

Creates a tensor of ones.

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

Creates a tensor filled with given value.

fn int_transpose( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Transposes an int tensor.

fn int_repeat_dim( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, times: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

fn int_not_equal( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: <Flex as BackendTypes>::IntTensorPrimitive, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise non-equality comparison.

fn int_not_equal_elem( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, out_dtype: BoolStore, ) -> <Flex as BackendTypes>::BoolTensorPrimitive

Element-wise non-equality comparison with a scalar.

fn int_sort( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, descending: bool, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

fn int_sort_with_indices( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, descending: bool, ) -> (<Flex as BackendTypes>::IntTensorPrimitive, <Flex as BackendTypes>::IntTensorPrimitive)

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

fn int_argsort( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, descending: bool, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

fn int_powi_scalar( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise power with a scalar.

fn int_powi_scalar_impl( lhs: <Flex as BackendTypes>::IntTensorPrimitive, rhs: Scalar, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Element-wise power with a scalar.

fn int_max_abs( tensor: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::IntTensorPrimitive

Gets the maximum absolute element in the tensor.

fn int_max_abs_dim( tensor: <Flex as BackendTypes>::IntTensorPrimitive, dim: usize, ) -> <Flex as BackendTypes>::IntTensorPrimitive

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

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

Creates a new tensor with values from the given range.

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

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

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: <Flex as BackendTypes>::FloatTensorPrimitive, weight: <Flex as BackendTypes>::FloatTensorPrimitive, bias: Option<<Flex as BackendTypes>::FloatTensorPrimitive>, options: ConvOptions<1>, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

One dimensional convolution.

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

Two dimensional convolution.

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

Two dimensional deformable convolution.

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

Backward pass for the deform_conv2d operation.

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

Three dimensional convolution.

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

One dimensional transposed convolution.

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

Two dimensional transposed convolution.

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

Three dimensional transposed convolution.

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

Two dimensional avg pooling.

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

Backward pass for the avg pooling 2d operation.

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

Two dimensional adaptive avg pooling.

fn adaptive_avg_pool2d_backward( x: <Flex as BackendTypes>::FloatTensorPrimitive, grad: <Flex as BackendTypes>::FloatTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Backward pass for the adaptive avg pooling 2d operation.

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

Two dimensional max pooling.

fn max_pool2d_with_indices( x: <Flex as BackendTypes>::FloatTensorPrimitive, 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: <Flex as BackendTypes>::FloatTensorPrimitive, _kernel_size: [usize; 2], _stride: [usize; 2], _padding: [usize; 2], _dilation: [usize; 2], _ceil_mode: bool, output_grad: <Flex as BackendTypes>::FloatTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> MaxPool2dBackward<Flex>

Backward pass for the max pooling 2d operation.

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

Down/up samples the input.

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

Backward pass for the interpolate operation.

fn attention( query: <Flex as BackendTypes>::FloatTensorPrimitive, key: <Flex as BackendTypes>::FloatTensorPrimitive, value: <Flex as BackendTypes>::FloatTensorPrimitive, mask: Option<<Flex as BackendTypes>::BoolTensorPrimitive>, attn_bias: Option<<Flex as BackendTypes>::FloatTensorPrimitive>, options: AttentionModuleOptions, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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: <Flex as BackendTypes>::FloatTensorPrimitive, dim: usize, n: Option<usize>, ) -> (<Flex as BackendTypes>::FloatTensorPrimitive, <Flex as BackendTypes>::FloatTensorPrimitive)

Real-valued FFT with optional size parameter.

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

Inverse real-valued FFT with optional output size.

fn embedding( weights: <Flex as BackendTypes>::FloatTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Embedding operation.

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

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

fn embedding_backward( weights: <Flex as BackendTypes>::FloatTensorPrimitive, output_grad: <Flex as BackendTypes>::FloatTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

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 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: &<Flex as BackendTypes>::Device, ) -> <Flex as BackendTypes>::QuantizedTensorPrimitive

Creates a new tensor from the data structure.

fn quantize_dynamic( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, scheme: &QuantScheme, ) -> <Flex as BackendTypes>::QuantizedTensorPrimitive

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

fn quantize( tensor: <Flex as BackendTypes>::FloatTensorPrimitive, scheme: &QuantScheme, qparams: QuantizationParametersPrimitive<Flex>, ) -> <Flex as BackendTypes>::QuantizedTensorPrimitive

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

fn dequantize( tensor: <Flex as BackendTypes>::QuantizedTensorPrimitive, dtype: FloatDType, ) -> <Flex as BackendTypes>::FloatTensorPrimitive

Convert the tensor back to a higher precision data type.

fn q_to_device( tensor: <Flex as BackendTypes>::QuantizedTensorPrimitive, _device: &<Flex as BackendTypes>::Device, ) -> <Flex as BackendTypes>::QuantizedTensorPrimitive

Moves the tensor to the given device.

fn q_reshape( tensor: <Flex as BackendTypes>::QuantizedTensorPrimitive, shape: Shape, ) -> <Flex as BackendTypes>::QuantizedTensorPrimitive

Reshapes a tensor.

async fn q_into_data( tensor: <Flex as BackendTypes>::QuantizedTensorPrimitive, ) -> Result<TensorData, ExecutionError>

Converts the tensor to a data structure.

fn q_swap_dims( tensor: <Flex as BackendTypes>::QuantizedTensorPrimitive, dim1: usize, dim2: usize, ) -> <Flex as BackendTypes>::QuantizedTensorPrimitive

Swaps two dimensions of a tensor.

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

Permutes the dimensions of a tensor.

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

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

fn q_expand( tensor: <Flex as BackendTypes>::QuantizedTensorPrimitive, shape: Shape, ) -> <Flex as BackendTypes>::QuantizedTensorPrimitive

Broadcasts the tensor to the given shape.

fn q_select( tensor: <Flex as BackendTypes>::QuantizedTensorPrimitive, dim: usize, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::QuantizedTensorPrimitive

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

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

Select tensor elements corresponding to the given slices.

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

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

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

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

fn q_gather( dim: usize, tensor: <Flex as BackendTypes>::QuantizedTensorPrimitive, indices: <Flex as BackendTypes>::IntTensorPrimitive, ) -> <Flex as BackendTypes>::QuantizedTensorPrimitive

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.