Struct burn_jit::JitBackend

pub struct JitBackend<R: Runtime> { /* private fields */ }

Generic tensor backend that can be compiled just-in-time to any shader runtime

Implementations

impl<R: Runtime> JitBackend<R>

pub fn new() -> Self

Constructs a new JitBackend.

Trait Implementations

impl<R: Runtime> ActivationOps<JitBackend<R>> for JitBackend<R>

fn leaky_relu<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, negative_slope: <B as Backend>::FloatElem ) -> <B as Backend>::FloatTensorPrimitive<D>

Applies the LeakyReLU activation function.

fn relu<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Applies the ReLU activation function.

fn relu_backward<const D: usize>( output: <B as Backend>::FloatTensorPrimitive<D>, grad: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Applies the ReLU activation function backward.

fn gelu<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Applies the Gelu activation function.

fn prelu<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, alpha: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Applies the PReLu activation function.

fn gelu_backward<const D: usize>( x: <B as Backend>::FloatTensorPrimitive<D>, grad: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Applies the Gelu activation function backward.

fn sigmoid<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Applies the Sigmoid activation function.

fn sigmoid_backward<const D: usize>( output: <B as Backend>::FloatTensorPrimitive<D>, grad: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Applies the Sigmoid activation function backward.

fn log_sigmoid<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Applies the LogSigmoid activation function.

fn log_sigmoid_backward<const D: usize>( x: <B as Backend>::FloatTensorPrimitive<D>, grad: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Applies the LogSigmoid activation function backward.

impl<R: Runtime> Backend for JitBackend<R>

type Device = <R as Runtime>::Device

Device type.

type FullPrecisionBridge = PrecisionBridge<<R as Runtime>::FullPrecisionRuntime>

A bridge that can cast tensors to full precision.

type FloatElem = <<R as Runtime>::Compiler as Compiler>::Float

Float element type.

type IntElem = <<R as Runtime>::Compiler as Compiler>::Int

Int element type.

type FloatTensorPrimitive<const D: usize> = JitTensor<R, <JitBackend<R> as Backend>::FloatElem, D>

Tensor primitive to be used for all float operations.

type IntTensorPrimitive<const D: usize> = JitTensor<R, <JitBackend<R> as Backend>::IntElem, D>

Tensor primitive to be used for all int operations.

type BoolTensorPrimitive<const D: usize> = JitTensor<R, u32, D>

Tensor primitive to be used for all bool operations.

fn name() -> String

Name of the backend.

fn seed(seed: u64)

Seed the backend.

fn ad_enabled() -> bool

If autodiff is enabled.

fn sync(device: &Self::Device)

Sync the backend, ensure that all computation are finished.

impl<ROrigin, RTarget> BackendBridge<JitBackend<ROrigin>> for PrecisionBridge<RTarget>

where ROrigin: Runtime, RTarget: Runtime<Device = ROrigin::Device, Server = ROrigin::Server, Channel = ROrigin::Channel>,

type Target = JitBackend<RTarget>

The target backend

fn into_target<const D: usize>( tensor: FloatTensor<JitBackend<ROrigin>, D>, device: Option<Device<Self::Target>> ) -> FloatTensor<Self::Target, D>

Transfer the tensor to the target backend.

fn from_target<const D: usize>( tensor: FloatTensor<Self::Target, D>, device: Option<Device<JitBackend<ROrigin>>> ) -> FloatTensor<JitBackend<ROrigin>, D>

Transfer the tensor from the target backend.

impl<R: Runtime> BoolTensorOps<JitBackend<R>> for JitBackend<R>

fn bool_empty<const D: usize>( shape: Shape<D>, device: &Device<Self> ) -> BoolTensor<Self, D>

Creates a new bool tensor.

fn bool_shape<const D: usize>(tensor: &BoolTensor<Self, D>) -> Shape<D>

Returns the shape of the tensor.

fn bool_into_data<const D: usize>( tensor: BoolTensor<Self, D> ) -> Reader<Data<bool, D>>

Converts the tensor to a data structure.

fn bool_from_data<const D: usize>( data: Data<bool, D>, device: &Device<Self> ) -> BoolTensor<Self, D>

Creates a tensor from the data structure.

fn bool_into_int<const D: usize>( tensor: BoolTensor<Self, D> ) -> IntTensor<Self, D>

Converts bool tensor to int tensor.

fn bool_device<const D: usize>(tensor: &BoolTensor<Self, D>) -> Device<Self>

Gets the device of the tensor.

fn bool_to_device<const D: usize>( tensor: BoolTensor<Self, D>, device: &Device<Self> ) -> BoolTensor<Self, D>

Moves the tensor to the device.

fn bool_reshape<const D1: usize, const D2: usize>( tensor: BoolTensor<Self, D1>, shape: Shape<D2> ) -> BoolTensor<Self, D2>

Reshapes the tensor.

fn bool_slice<const D1: usize, const D2: usize>( tensor: BoolTensor<Self, D1>, ranges: [Range<usize>; D2] ) -> BoolTensor<Self, D1>

Gets the values from the tensor for the given ranges.

fn bool_slice_assign<const D1: usize, const D2: usize>( tensor: BoolTensor<Self, D1>, ranges: [Range<usize>; D2], value: BoolTensor<Self, D1> ) -> BoolTensor<Self, D1>

Sets the values in the tensor for the given ranges.

fn bool_equal<const D: usize>( lhs: BoolTensor<Self, D>, rhs: BoolTensor<Self, D> ) -> BoolTensor<Self, D>

Equates the two tensors.

fn bool_not<const D: usize>(tensor: BoolTensor<Self, D>) -> BoolTensor<Self, D>

Inverses boolean values.

fn bool_into_float<const D: usize>( tensor: BoolTensor<Self, D> ) -> FloatTensor<Self, D>

Converts bool tensor to float tensor.

fn bool_swap_dims<const D: usize>( tensor: BoolTensor<Self, D>, dim1: usize, dim2: usize ) -> BoolTensor<Self, D>

Swaps two dimensions of a bool tensor.

fn bool_repeat<const D: usize>( tensor: BoolTensor<Self, D>, dim: usize, times: usize ) -> BoolTensor<Self, D>

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

fn bool_permute<const D: usize>( tensor: BoolTensor<Self, D>, axes: [usize; D] ) -> BoolTensor<Self, D>

Permutes the dimensions of a tensor.

fn bool_expand<const D1: usize, const D2: usize>( tensor: BoolTensor<Self, D1>, shape: Shape<D2> ) -> BoolTensor<Self, D2>

Broadcasts the bool tensor to the given shape.

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

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

fn bool_to_data<const D: usize>( tensor: &<B as Backend>::BoolTensorPrimitive<D> ) -> Reader<Data<bool, D>>

Gets the data from the tensor.

fn bool_cat<const D: usize>( tensors: Vec<<B as Backend>::BoolTensorPrimitive<D>>, dim: usize ) -> <B as Backend>::BoolTensorPrimitive<D>

Concatenates the tensors along the given dimension.

fn bool_not_equal<const D: usize>( lhs: <B as Backend>::BoolTensorPrimitive<D>, rhs: <B as Backend>::BoolTensorPrimitive<D> ) -> <B as Backend>::BoolTensorPrimitive<D>

Element-wise non-equality comparison.

fn bool_transpose<const D: usize>( tensor: <B as Backend>::BoolTensorPrimitive<D> ) -> <B as Backend>::BoolTensorPrimitive<D>

Transposes a bool tensor.

fn bool_narrow<const D: usize>( tensor: <B as Backend>::BoolTensorPrimitive<D>, dim: usize, start: usize, length: usize ) -> <B as Backend>::BoolTensorPrimitive<D>

Returns a new tensor with the given dimension narrowed to the given range.

fn bool_chunk<const D: usize>( tensor: <B as Backend>::BoolTensorPrimitive<D>, chunks: usize, dim: usize ) -> Vec<<B as Backend>::BoolTensorPrimitive<D>>

Split the tensor along the given dimension into chunks.

fn bool_any<const D: usize>( tensor: <B as Backend>::BoolTensorPrimitive<D> ) -> <B as Backend>::BoolTensorPrimitive<1>

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

fn bool_any_dim<const D: usize>( tensor: <B as Backend>::BoolTensorPrimitive<D>, dim: usize ) -> <B as Backend>::BoolTensorPrimitive<D>

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

fn bool_all<const D: usize>( tensor: <B as Backend>::BoolTensorPrimitive<D> ) -> <B as Backend>::BoolTensorPrimitive<1>

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

fn bool_all_dim<const D: usize>( tensor: <B as Backend>::BoolTensorPrimitive<D>, dim: usize ) -> <B as Backend>::BoolTensorPrimitive<D>

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

fn bool_argwhere<const D: usize>( tensor: <B as Backend>::BoolTensorPrimitive<D> ) -> <B as Backend>::IntTensorPrimitive<2>

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

fn bool_nonzero<const D: usize>( tensor: <B as Backend>::BoolTensorPrimitive<D> ) -> Vec<<B as Backend>::IntTensorPrimitive<1>>

Compute the indices of the elements that are non-zero.

impl<R: Runtime> Clone for JitBackend<R>

fn clone(&self) -> Self

Returns a copy of the value.

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

Performs copy-assignment from source.

impl<R: Runtime> Debug for JitBackend<R>

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

Formats the value using the given formatter.

impl<R: Runtime> Default for JitBackend<R>

fn default() -> Self

Returns the “default value” for a type.

impl<R: Runtime> FloatTensorOps<JitBackend<R>> for JitBackend<R>

fn float_from_data<const D: usize>( data: Data<FloatElem<Self>, D>, device: &Device<Self> ) -> FloatTensor<Self, D>

Creates a new tensor from the data structure.

fn float_random<const D: usize>( shape: Shape<D>, distribution: Distribution, device: &Device<Self> ) -> FloatTensor<Self, D>

Creates a new tensor with random values.

fn float_shape<const D: usize>(tensor: &FloatTensor<Self, D>) -> Shape<D>

Gets the shape of the tensor.

fn float_into_data<const D: usize>( tensor: FloatTensor<Self, D> ) -> Reader<Data<FloatElem<Self>, D>>

Converts the tensor to a data structure.

fn float_device<const D: usize>(tensor: &FloatTensor<Self, D>) -> Device<Self>

Gets the device of the tensor.

fn float_to_device<const D: usize>( tensor: FloatTensor<Self, D>, device: &Device<Self> ) -> FloatTensor<Self, D>

Moves the tensor to the given device.

fn float_empty<const D: usize>( shape: Shape<D>, device: &Device<Self> ) -> FloatTensor<Self, D>

Creates an empty tensor with the given shape.

fn float_add<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Adds two tensors together.

fn float_add_scalar<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatElem<Self> ) -> FloatTensor<Self, D>

Adds a scalar to a tensor.

fn float_zeros<const D: usize>( shape: Shape<D>, device: &Device<Self> ) -> FloatTensor<Self, D>

Creates a new tensor with zeros.

fn float_full<const D: usize>( shape: Shape<D>, fill_value: FloatElem<Self>, device: &R::Device ) -> FloatTensor<Self, D>

Creates a tensor filled with given value.

fn float_ones<const D: usize>( shape: Shape<D>, device: &Device<Self> ) -> FloatTensor<Self, D>

Creates a new tensor with ones.

fn float_sub<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Subtracts two tensors.

fn float_sub_scalar<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatElem<Self> ) -> FloatTensor<Self, D>

Subtracts a scalar from a tensor.

fn float_mul<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Multiplies two tensors together element-wise.

fn float_mul_scalar<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatElem<Self> ) -> FloatTensor<Self, D>

Multiplies a tensor by a scalar.

fn float_div<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Divides two tensors element-wise.

fn float_div_scalar<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatElem<Self> ) -> FloatTensor<Self, D>

Divides a tensor by a scalar.

fn float_matmul<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Multiplies two tensors together using matrix multiplication.

fn float_swap_dims<const D: usize>( tensor: FloatTensor<Self, D>, dim1: usize, dim2: usize ) -> FloatTensor<Self, D>

Swaps two dimensions of a tensor.

fn float_reshape<const D1: usize, const D2: usize>( tensor: FloatTensor<Self, D1>, shape: Shape<D2> ) -> FloatTensor<Self, D2>

Reshapes a tensor.

fn float_gather<const D: usize>( dim: usize, tensor: FloatTensor<Self, D>, indices: IntTensor<Self, D> ) -> FloatTensor<Self, D>

Gather elements from a tensor.

fn float_scatter<const D: usize>( dim: usize, tensor: FloatTensor<Self, D>, indices: IntTensor<Self, D>, value: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Scatter elements into a tensor.

fn float_select<const D: usize>( tensor: FloatTensor<Self, D>, dim: usize, indices: IntTensor<Self, 1> ) -> FloatTensor<Self, D>

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

fn float_select_assign<const D: usize>( tensor: FloatTensor<Self, D>, dim: usize, indices: IntTensor<Self, 1>, value: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

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

fn float_slice<const D1: usize, const D2: usize>( tensor: FloatTensor<Self, D1>, ranges: [Range<usize>; D2] ) -> FloatTensor<Self, D1>

Select tensor elements corresponding for the given ranges.

fn float_slice_assign<const D1: usize, const D2: usize>( tensor: FloatTensor<Self, D1>, ranges: [Range<usize>; D2], value: FloatTensor<Self, D1> ) -> FloatTensor<Self, D1>

Assign the selected elements corresponding for the given ranges to the given value.

fn float_mask_where<const D: usize>( tensor: FloatTensor<Self, D>, mask: BoolTensor<Self, D>, value: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

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

fn float_mask_fill<const D: usize>( tensor: FloatTensor<Self, D>, mask: BoolTensor<Self, D>, value: FloatElem<Self> ) -> FloatTensor<Self, D>

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

fn float_equal<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatTensor<Self, D> ) -> BoolTensor<Self, D>

Equal comparison of two tensors.

fn float_equal_elem<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatElem<Self> ) -> BoolTensor<Self, D>

Equal comparison of a tensor and a scalar.

fn float_greater<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatTensor<Self, D> ) -> BoolTensor<Self, D>

Greater than comparison of two tensors.

fn float_greater_elem<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatElem<Self> ) -> BoolTensor<Self, D>

Greater than comparison of a tensor and a scalar.

fn float_greater_equal<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatTensor<Self, D> ) -> BoolTensor<Self, D>

Greater than or equal comparison of two tensors.

fn float_greater_equal_elem<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatElem<Self> ) -> BoolTensor<Self, D>

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

fn float_lower<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatTensor<Self, D> ) -> BoolTensor<Self, D>

Less than comparison of two tensors.

fn float_lower_elem<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatElem<Self> ) -> BoolTensor<Self, D>

Less than comparison of a tensor and a scalar.

fn float_lower_equal<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatTensor<Self, D> ) -> BoolTensor<Self, D>

Less than or equal comparison of two tensors.

fn float_lower_equal_elem<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatElem<Self> ) -> BoolTensor<Self, D>

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

fn float_sum<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, 1>

Sum of all elements in a tensor.

fn float_sum_dim<const D: usize>( tensor: FloatTensor<Self, D>, dim: usize ) -> FloatTensor<Self, D>

Sum of all elements in a tensor along a dimension.

fn float_mean_dim<const D: usize>( tensor: FloatTensor<Self, D>, dim: usize ) -> FloatTensor<Self, D>

Mean of all elements in a tensor along a dimension.

fn float_prod<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, 1>

Product of all elements in a tensor.

fn float_prod_dim<const D: usize>( tensor: FloatTensor<Self, D>, dim: usize ) -> FloatTensor<Self, D>

Product of all elements in a tensor along a dimension.

fn float_exp<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Returns a new tensor with exponential values.

fn float_log<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Returns a new tensor with natural logarithm values.

fn float_log1p<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

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

fn float_powf_scalar<const D: usize>( lhs: FloatTensor<Self, D>, rhs: f32 ) -> FloatTensor<Self, D>

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

fn float_sqrt<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Returns a new tensor with square root values.

fn float_abs<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Returns a new tensor with absolute values.

fn float_cos<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Returns a new tensor with cosine values.

fn float_sin<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Returns a new tensor with sine values.

fn float_tanh<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Returns a new tensor with tangent values.

fn float_erf<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Returns a new tensor with the error function values.

fn float_argmax<const D: usize>( tensor: FloatTensor<Self, D>, dim: usize ) -> IntTensor<Self, D>

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

fn float_argmin<const D: usize>( tensor: FloatTensor<Self, D>, dim: usize ) -> IntTensor<Self, D>

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

fn float_into_int<const D: usize>( tensor: FloatTensor<Self, D> ) -> IntTensor<Self, D>

Converts float tensor to int tensor.

fn float_clamp<const D: usize>( tensor: FloatTensor<Self, D>, min: FloatElem<Self>, max: FloatElem<Self> ) -> FloatTensor<Self, D>

Clamps a tensor between a minimum and maximum value.

fn float_recip<const D: usize>( tensor: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Calculates the reciprocals element-wise

fn float_repeat<const D: usize>( tensor: FloatTensor<Self, D>, dim: usize, times: usize ) -> FloatTensor<Self, D>

Repeat the tensor along the given dimension.

fn float_powf<const D: usize>( lhs: FloatTensor<Self, D>, rhs: FloatTensor<Self, D> ) -> FloatTensor<Self, D>

Element-wise power with a FloatTensor.

fn float_permute<const D: usize>( tensor: FloatTensor<Self, D>, axes: [usize; D] ) -> FloatTensor<Self, D>

Permutes the dimensions of a tensor.

fn float_expand<const D1: usize, const D2: usize>( tensor: FloatTensor<Self, D1>, shape: Shape<D2> ) -> FloatTensor<Self, D2>

Broadcasts the float tensor to the given shape.

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

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

fn float_to_data<const D: usize>( tensor: &<B as Backend>::FloatTensorPrimitive<D> ) -> Reader<Data<<B as Backend>::FloatElem, D>>

Converts the tensor to a data structure.

fn float_clamp_min<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, min: <B as Backend>::FloatElem ) -> <B as Backend>::FloatTensorPrimitive<D>

Clamps a tensor under a minimum value.

fn float_clamp_max<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, max: <B as Backend>::FloatElem ) -> <B as Backend>::FloatTensorPrimitive<D>

Clamps a tensor over a maximum value.

fn float_neg<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Negates a tensor element-wise.

fn float_transpose<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Transposes a tensor.

fn float_not_equal<const D: usize>( lhs: <B as Backend>::FloatTensorPrimitive<D>, rhs: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::BoolTensorPrimitive<D>

Element-wise non-equality comparison.

fn float_not_equal_elem<const D: usize>( lhs: <B as Backend>::FloatTensorPrimitive<D>, rhs: <B as Backend>::FloatElem ) -> <B as Backend>::BoolTensorPrimitive<D>

Element-wise non-equality comparison with a scalar.

fn float_detach<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Detaches a tensor from the computation graph.

fn float_set_require_grad<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, _require_grad: bool ) -> <B as Backend>::FloatTensorPrimitive<D>

Sets the require_grad flag of a tensor.

fn float_is_require_grad<const D: usize>( _tensor: &<B as Backend>::FloatTensorPrimitive<D> ) -> bool

Returns the require_grad flag of a tensor.

fn float_mean<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<1>

Mean of all elements in a tensor.

fn float_into_full_precision<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <<<B as Backend>::FullPrecisionBridge as BackendBridge<B>>::Target as Backend>::FloatTensorPrimitive<D>

Converts a tensor to full precision.

fn float_from_full_precision<const D: usize>( tensor: <<<B as Backend>::FullPrecisionBridge as BackendBridge<B>>::Target as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Converts a tensor from full precision.

fn float_powi<const D: usize>( lhs: <B as Backend>::FloatTensorPrimitive<D>, rhs: <B as Backend>::IntTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Element-wise power with an IntTensor.

fn float_powi_scalar<const D: usize>( lhs: <B as Backend>::FloatTensorPrimitive<D>, rhs: <B as Backend>::IntElem ) -> <B as Backend>::FloatTensorPrimitive<D>

raises a tensor to the power of an int scalar.

fn float_cat<const D: usize>( tensors: Vec<<B as Backend>::FloatTensorPrimitive<D>>, dim: usize ) -> <B as Backend>::FloatTensorPrimitive<D>

Concatenates tensors along a dimension.

fn float_max<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<1>

Gets the maximum element of a tensor.

fn float_max_dim<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, dim: usize ) -> <B as Backend>::FloatTensorPrimitive<D>

Gets the maximum elements of a tensor along an axis.

fn float_max_dim_with_indices<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, dim: usize ) -> (<B as Backend>::FloatTensorPrimitive<D>, <B as Backend>::IntTensorPrimitive<D>)

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

fn float_min<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<1>

Gets the minimum element of a tensor.

fn float_min_dim<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, dim: usize ) -> <B as Backend>::FloatTensorPrimitive<D>

Gets the minimum elements of a tensor along an axis.

fn float_min_dim_with_indices<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, dim: usize ) -> (<B as Backend>::FloatTensorPrimitive<D>, <B as Backend>::IntTensorPrimitive<D>)

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

fn float_narrow<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, dim: usize, start: usize, length: usize ) -> <B as Backend>::FloatTensorPrimitive<D>

Returns a new tensor with the given dimension narrowed to the given range.

fn float_chunk<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, chunks: usize, dim: usize ) -> Vec<<B as Backend>::FloatTensorPrimitive<D>>

Split the tensor along the given dimension into chunks.

fn float_any<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::BoolTensorPrimitive<1>

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

fn float_any_dim<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, dim: usize ) -> <B as Backend>::BoolTensorPrimitive<D>

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

fn float_all<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::BoolTensorPrimitive<1>

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

fn float_all_dim<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, dim: usize ) -> <B as Backend>::BoolTensorPrimitive<D>

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

fn float_sign<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::FloatTensorPrimitive<D>

Returns the signs of the float tensor.

fn float_sort<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, dim: usize, descending: bool ) -> <B as Backend>::FloatTensorPrimitive<D>

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

fn float_sort_with_indices<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, dim: usize, descending: bool ) -> (<B as Backend>::FloatTensorPrimitive<D>, <B as Backend>::IntTensorPrimitive<D>)

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

fn float_argsort<const D: usize>( tensor: <B as Backend>::FloatTensorPrimitive<D>, dim: usize, descending: bool ) -> <B as Backend>::IntTensorPrimitive<D>

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

impl<R: Runtime> FusionBackend for JitBackend<R>

type OptimizationState = JitOptimizationState

The state that can be serialized for an optimization.

type Optimization = JitOptimization<R>

Optimization type for the backend.

type FusionDevice = <R as Runtime>::Device

The device type that can return an ID.

type Handle = JitFusionHandle<R>

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

type FusionClient = MutexFusionClient<JitBackend<R>>

What kind of client should be used.

fn optimizations( device: R::Device ) -> Vec<Box<dyn OptimizationBuilder<Self::Optimization>>>

The list of optimizations that will be used to optimize the computational graph.

fn float_tensor<const D: usize>( handle: Self::Handle, shape: Shape<D> ) -> Self::FloatTensorPrimitive<D>

Convert a handle to a float tensor.

fn int_tensor<const D: usize>( handle: Self::Handle, shape: Shape<D> ) -> Self::IntTensorPrimitive<D>

Convert a handle to an int tensor.

fn bool_tensor<const D: usize>( handle: Self::Handle, shape: Shape<D> ) -> Self::BoolTensorPrimitive<D>

Convert a handle to a bool tensor.

fn float_tensor_handle<const D: usize>( tensor: Self::FloatTensorPrimitive<D> ) -> Self::Handle

Convert a float tensor to a handle.

fn int_tensor_handle<const D: usize>( tensor: Self::IntTensorPrimitive<D> ) -> Self::Handle

Convert an int tensor to a handle.

fn bool_tensor_handle<const D: usize>( tensor: Self::BoolTensorPrimitive<D> ) -> Self::Handle

Convert a bool tensor to a handle.

impl<R: Runtime> IntTensorOps<JitBackend<R>> for JitBackend<R>

fn int_empty<const D: usize>( shape: Shape<D>, device: &Device<Self> ) -> IntTensor<Self, D>

Creates a new int tensor.

fn int_shape<const D: usize>(tensor: &IntTensor<Self, D>) -> Shape<D>

Returns the shape of the tensor.

fn int_into_data<const D: usize>( tensor: IntTensor<Self, D> ) -> Reader<Data<IntElem<Self>, D>>

Converts the tensor to a data structure.

fn int_from_data<const D: usize>( data: Data<IntElem<Self>, D>, device: &Device<Self> ) -> IntTensor<Self, D>

Creates a tensor from the data structure.

fn int_device<const D: usize>(tensor: &IntTensor<Self, D>) -> Device<Self>

Gets the device of the tensor.

fn int_to_device<const D: usize>( tensor: IntTensor<Self, D>, device: &Device<Self> ) -> IntTensor<Self, D>

Moves the tensor to the given device.

fn int_reshape<const D1: usize, const D2: usize>( tensor: IntTensor<Self, D1>, shape: Shape<D2> ) -> IntTensor<Self, D2>

Reshapes the tensor.

fn int_slice<const D1: usize, const D2: usize>( tensor: IntTensor<Self, D1>, ranges: [Range<usize>; D2] ) -> IntTensor<Self, D1>

Gets the element at the given indices.

fn int_slice_assign<const D1: usize, const D2: usize>( tensor: IntTensor<Self, D1>, ranges: [Range<usize>; D2], value: IntTensor<Self, D1> ) -> IntTensor<Self, D1>

Sets the element at the given indices.

fn int_mask_where<const D: usize>( tensor: IntTensor<Self, D>, mask: BoolTensor<Self, D>, value: IntTensor<Self, D> ) -> IntTensor<Self, D>

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

fn int_mask_fill<const D: usize>( tensor: IntTensor<Self, D>, mask: BoolTensor<Self, D>, value: IntElem<Self> ) -> IntTensor<Self, D>

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

fn int_gather<const D: usize>( dim: usize, tensor: IntTensor<Self, D>, indices: IntTensor<Self, D> ) -> IntTensor<Self, D>

Gather elements from the tensor at the given indices.

fn int_scatter<const D: usize>( dim: usize, tensor: IntTensor<Self, D>, indices: IntTensor<Self, D>, value: IntTensor<Self, D> ) -> IntTensor<Self, D>

Scatter a given value to the tensor at the given indices.

fn int_select<const D: usize>( tensor: IntTensor<Self, D>, dim: usize, indices: IntTensor<Self, 1> ) -> IntTensor<Self, D>

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

fn int_select_assign<const D: usize>( tensor: IntTensor<Self, D>, dim: usize, indices: IntTensor<Self, 1>, value: IntTensor<Self, D> ) -> IntTensor<Self, D>

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

fn int_equal<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntTensor<Self, D> ) -> BoolTensor<Self, D>

Element-wise equality comparison.

fn int_equal_elem<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntElem<Self> ) -> BoolTensor<Self, D>

Element-wise equality comparison with a scalar.

fn int_greater<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntTensor<Self, D> ) -> BoolTensor<Self, D>

Element-wise greater than comparison.

fn int_greater_elem<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntElem<Self> ) -> BoolTensor<Self, D>

Element-wise greater than comparison with a scalar.

fn int_greater_equal<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntTensor<Self, D> ) -> BoolTensor<Self, D>

Element-wise greater than or equal comparison.

fn int_greater_equal_elem<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntElem<Self> ) -> BoolTensor<Self, D>

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

fn int_lower<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntTensor<Self, D> ) -> BoolTensor<Self, D>

Element-wise less than comparison.

fn int_lower_elem<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntElem<Self> ) -> BoolTensor<Self, D>

Element-wise less than comparison with a scalar.

fn int_lower_equal<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntTensor<Self, D> ) -> BoolTensor<Self, D>

Element-wise less than or equal comparison.

fn int_lower_equal_elem<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntElem<Self> ) -> BoolTensor<Self, D>

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

fn int_add<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntTensor<Self, D> ) -> IntTensor<Self, D>

Element-wise addition.

fn int_add_scalar<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntElem<Self> ) -> IntTensor<Self, D>

Element-wise addition with a scalar.

fn int_sub<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntTensor<Self, D> ) -> IntTensor<Self, D>

Element-wise subtraction.

fn int_sub_scalar<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntElem<Self> ) -> IntTensor<Self, D>

Element-wise subtraction with a scalar.

fn int_mul<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntTensor<Self, D> ) -> IntTensor<Self, D>

Element-wise multiplication.

fn int_mul_scalar<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntElem<Self> ) -> IntTensor<Self, D>

Element-wise multiplication with a scalar.

fn int_div<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntTensor<Self, D> ) -> IntTensor<Self, D>

Element-wise division.

fn int_div_scalar<const D: usize>( lhs: IntTensor<Self, D>, rhs: IntElem<Self> ) -> IntTensor<Self, D>

Element-wise division with a scalar.

fn int_zeros<const D: usize>( shape: Shape<D>, device: &Device<Self> ) -> IntTensor<Self, D>

Creates a tensor of zeros.

fn int_ones<const D: usize>( shape: Shape<D>, device: &Device<Self> ) -> IntTensor<Self, D>

Creates a tensor of ones.

fn int_sum<const D: usize>(tensor: IntTensor<Self, D>) -> IntTensor<Self, 1>

Sums all elements in the tensor.

fn int_sum_dim<const D: usize>( tensor: IntTensor<Self, D>, dim: usize ) -> IntTensor<Self, D>

Sums all elements in the tensor along a dimension.

fn int_prod<const D: usize>(tensor: IntTensor<Self, D>) -> IntTensor<Self, 1>

Computes the product of all elements in the tensor.

fn int_prod_dim<const D: usize>( tensor: IntTensor<Self, D>, dim: usize ) -> IntTensor<Self, D>

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

fn int_mean_dim<const D: usize>( tensor: IntTensor<Self, D>, dim: usize ) -> IntTensor<Self, D>

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

fn int_argmax<const D: usize>( tensor: IntTensor<Self, D>, dim: usize ) -> IntTensor<Self, D>

Gets the indices of the maximum elements along a dimension.

fn int_argmin<const D: usize>( tensor: IntTensor<Self, D>, dim: usize ) -> IntTensor<Self, D>

Gets the indices of the minimum elements along a dimension.

fn int_clamp<const D: usize>( tensor: IntTensor<Self, D>, min: IntElem<Self>, max: IntElem<Self> ) -> IntTensor<Self, D>

Clamps a tensor between a minimum and maximum value.

fn int_abs<const D: usize>(tensor: IntTensor<Self, D>) -> IntTensor<Self, D>

Returns a new tensor with absolute values.

fn int_into_float<const D: usize>( tensor: IntTensor<Self, D> ) -> FloatTensor<Self, D>

Converts int tensor to float tensor.

fn int_swap_dims<const D: usize>( tensor: IntTensor<Self, D>, dim1: usize, dim2: usize ) -> IntTensor<Self, D>

Swaps two dimensions of an int tensor.

fn int_repeat<const D: usize>( tensor: IntTensor<Self, D>, dim: usize, times: usize ) -> IntTensor<Self, D>

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

fn int_random<const D: usize>( shape: Shape<D>, distribution: Distribution, device: &Device<Self> ) -> IntTensor<Self, D>

Creates a new int tensor with random values.

fn int_permute<const D: usize>( tensor: IntTensor<Self, D>, axes: [usize; D] ) -> IntTensor<Self, D>

Permutes the dimensions of a tensor.

fn int_expand<const D1: usize, const D2: usize>( tensor: IntTensor<Self, D1>, shape: Shape<D2> ) -> IntTensor<Self, D2>

Broadcasts the int tensor to the given shape.

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

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

fn int_to_data<const D: usize>( tensor: &<B as Backend>::IntTensorPrimitive<D> ) -> Reader<Data<<B as Backend>::IntElem, D>>

Gets the data from the tensor.

fn int_cat<const D: usize>( tensors: Vec<<B as Backend>::IntTensorPrimitive<D>>, dim: usize ) -> <B as Backend>::IntTensorPrimitive<D>

Concatenates the given tensors along the given dimension.

fn int_not_equal<const D: usize>( lhs: <B as Backend>::IntTensorPrimitive<D>, rhs: <B as Backend>::IntTensorPrimitive<D> ) -> <B as Backend>::BoolTensorPrimitive<D>

Element-wise non-equality comparison.

fn int_not_equal_elem<const D: usize>( lhs: <B as Backend>::IntTensorPrimitive<D>, rhs: <B as Backend>::IntElem ) -> <B as Backend>::BoolTensorPrimitive<D>

Element-wise non-equality comparison with a scalar.

fn int_powi<const D: usize>( lhs: <B as Backend>::IntTensorPrimitive<D>, rhs: <B as Backend>::IntTensorPrimitive<D> ) -> <B as Backend>::IntTensorPrimitive<D>

Element-wise power with a IntTensor.

fn int_powf<const D: usize>( lhs: <B as Backend>::IntTensorPrimitive<D>, rhs: <B as Backend>::FloatTensorPrimitive<D> ) -> <B as Backend>::IntTensorPrimitive<D>

Element-wise power with a floatTensor.

fn int_powi_scalar<const D: usize>( lhs: <B as Backend>::IntTensorPrimitive<D>, rhs: <B as Backend>::IntElem ) -> <B as Backend>::IntTensorPrimitive<D>

Element-wise power with a scalar.

fn int_powf_scalar<const D: usize>( lhs: <B as Backend>::IntTensorPrimitive<D>, rhs: f32 ) -> <B as Backend>::IntTensorPrimitive<D>

Element-wise power with a floatTensor.

fn int_clamp_min<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, min: <B as Backend>::IntElem ) -> <B as Backend>::IntTensorPrimitive<D>

Clamps a tensor under a minimum value.

fn int_clamp_max<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, max: <B as Backend>::IntElem ) -> <B as Backend>::IntTensorPrimitive<D>

Clamps a tensor over a maximum value.

fn int_neg<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D> ) -> <B as Backend>::IntTensorPrimitive<D>

Element-wise negation.

fn int_full<const D: usize>( shape: Shape<D>, fill_value: <B as Backend>::IntElem, device: &<B as Backend>::Device ) -> <B as Backend>::IntTensorPrimitive<D>

Creates a tensor filled with given value.

fn int_mean<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D> ) -> <B as Backend>::IntTensorPrimitive<1>

Computes the mean of all elements in the tensor.

fn int_max<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D> ) -> <B as Backend>::IntTensorPrimitive<1>

Gets the maximum element in the tensor.

fn int_max_dim<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, dim: usize ) -> <B as Backend>::IntTensorPrimitive<D>

Gets the maximum element in the tensor along a dimension.

fn int_max_dim_with_indices<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, dim: usize ) -> (<B as Backend>::IntTensorPrimitive<D>, <B as Backend>::IntTensorPrimitive<D>)

Gets the maximum elements and corresponding indices along a dimension.

fn int_min<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D> ) -> <B as Backend>::IntTensorPrimitive<1>

Gets the minimum element in the tensor.

fn int_min_dim<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, dim: usize ) -> <B as Backend>::IntTensorPrimitive<D>

Gets the minimum elements in the tensor along a dimension.

fn int_min_dim_with_indices<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, dim: usize ) -> (<B as Backend>::IntTensorPrimitive<D>, <B as Backend>::IntTensorPrimitive<D>)

Gets the minimum elements and corresponding indices along a dimension.

fn int_transpose<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D> ) -> <B as Backend>::IntTensorPrimitive<D>

Transposes an int tensor.

fn int_narrow<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, dim: usize, start: usize, length: usize ) -> <B as Backend>::IntTensorPrimitive<D>

Returns a new tensor with the given dimension narrowed to the given range.

fn int_chunk<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, chunks: usize, dim: usize ) -> Vec<<B as Backend>::IntTensorPrimitive<D>>

Split the tensor along the given dimension into chunks.

fn int_arange_step( range: Range<i64>, step: usize, device: &<B as Backend>::Device ) -> <B as Backend>::IntTensorPrimitive<1>

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

fn int_arange( range: Range<i64>, device: &<B as Backend>::Device ) -> <B as Backend>::IntTensorPrimitive<1>

Creates a new tensor with values from the given range.

fn int_any<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D> ) -> <B as Backend>::BoolTensorPrimitive<1>

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

fn int_any_dim<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, dim: usize ) -> <B as Backend>::BoolTensorPrimitive<D>

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

fn int_all<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D> ) -> <B as Backend>::BoolTensorPrimitive<1>

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

fn int_all_dim<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, dim: usize ) -> <B as Backend>::BoolTensorPrimitive<D>

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

fn int_sign<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D> ) -> <B as Backend>::IntTensorPrimitive<D>

Returns the signs of the int tensor.

fn int_sort<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, dim: usize, descending: bool ) -> <B as Backend>::IntTensorPrimitive<D>

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

fn int_sort_with_indices<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, dim: usize, descending: bool ) -> (<B as Backend>::IntTensorPrimitive<D>, <B as Backend>::IntTensorPrimitive<D>)

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

fn int_argsort<const D: usize>( tensor: <B as Backend>::IntTensorPrimitive<D>, dim: usize, descending: bool ) -> <B as Backend>::IntTensorPrimitive<D>

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

impl<R: Runtime> ModuleOps<JitBackend<R>> for JitBackend<R>

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

Two dimensional convolution.

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

Two dimensional transposed convolution.

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

Two dimensional avg pooling.

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

Backward pass for the avg pooling 2d operation.

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

Two dimensional max pooling.

fn max_pool2d_with_indices( x: FloatTensor<Self, 4>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2] ) -> MaxPool2dWithIndices<Self>

Two dimensional max pooling with indices.

fn max_pool2d_with_indices_backward( x: FloatTensor<Self, 4>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], output_grad: FloatTensor<Self, 4>, indices: IntTensor<Self, 4> ) -> MaxPool2dBackward<Self>

Backward pass for the max pooling 2d operation.

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

Two dimensional adaptive avg pooling.

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

Backward pass for the adaptive avg pooling 2d operation.

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

Down/up samples the input.

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

Backward pass for the interpolate operation.

fn embedding( weights: <B as Backend>::FloatTensorPrimitive<2>, indices: <B as Backend>::IntTensorPrimitive<2> ) -> <B as Backend>::FloatTensorPrimitive<3>

Embedding operation.

fn embedding_backward( weights: <B as Backend>::FloatTensorPrimitive<2>, output_grad: <B as Backend>::FloatTensorPrimitive<3>, indices: <B as Backend>::IntTensorPrimitive<2> ) -> <B as Backend>::FloatTensorPrimitive<2>

Embedding backward operation.

fn conv1d( x: <B as Backend>::FloatTensorPrimitive<3>, weight: <B as Backend>::FloatTensorPrimitive<3>, bias: Option<<B as Backend>::FloatTensorPrimitive<1>>, options: ConvOptions<1> ) -> <B as Backend>::FloatTensorPrimitive<3>

One dimensional convolution.

fn conv1d_backward( x: <B as Backend>::FloatTensorPrimitive<3>, weight: <B as Backend>::FloatTensorPrimitive<3>, bias: Option<<B as Backend>::FloatTensorPrimitive<1>>, output_grad: <B as Backend>::FloatTensorPrimitive<3>, options: ConvOptions<1> ) -> Conv1dBackward<B>

Backward pass for the conv1d operation.

fn conv2d_backward( x: <B as Backend>::FloatTensorPrimitive<4>, weight: <B as Backend>::FloatTensorPrimitive<4>, bias: Option<<B as Backend>::FloatTensorPrimitive<1>>, output_grad: <B as Backend>::FloatTensorPrimitive<4>, options: ConvOptions<2> ) -> Conv2dBackward<B>

Backward pass for the conv2d operation.

fn conv_transpose1d( x: <B as Backend>::FloatTensorPrimitive<3>, weight: <B as Backend>::FloatTensorPrimitive<3>, bias: Option<<B as Backend>::FloatTensorPrimitive<1>>, options: ConvTransposeOptions<1> ) -> <B as Backend>::FloatTensorPrimitive<3>

One dimensional transposed convolution.

fn conv_transpose1d_backward( x: <B as Backend>::FloatTensorPrimitive<3>, weight: <B as Backend>::FloatTensorPrimitive<3>, bias: Option<<B as Backend>::FloatTensorPrimitive<1>>, output_grad: <B as Backend>::FloatTensorPrimitive<3>, options: ConvTransposeOptions<1> ) -> Conv1dBackward<B>

Backward pass for the conv transpose 1d operation.

fn conv_transpose2d_backward( x: <B as Backend>::FloatTensorPrimitive<4>, weight: <B as Backend>::FloatTensorPrimitive<4>, bias: Option<<B as Backend>::FloatTensorPrimitive<1>>, output_grad: <B as Backend>::FloatTensorPrimitive<4>, options: ConvTransposeOptions<2> ) -> Conv2dBackward<B>

Backward pass for the conv transpose 2d operation.

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

Four-dimensional unfolding.

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

One dimensional avg pooling.

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

Backward pass for the avg pooling 1d operation.

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

One dimensional adaptive avg pooling.

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

Backward pass for the adaptive avg pooling 1d operation.

fn max_pool1d( x: <B as Backend>::FloatTensorPrimitive<3>, kernel_size: usize, stride: usize, padding: usize, dilation: usize ) -> <B as Backend>::FloatTensorPrimitive<3>

One dimensional max pooling.

fn max_pool1d_with_indices( x: <B as Backend>::FloatTensorPrimitive<3>, kernel_size: usize, stride: usize, padding: usize, dilation: usize ) -> MaxPool1dWithIndices<B>

One dimensional max pooling with indices.

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

Backward pass for the max pooling 1d operation.