Trait AutodiffBackend 

pub trait AutodiffBackend: Backend {
    type InnerBackend: Backend<Device = Self::Device>;
    type Gradients: Send;

    // Required methods
    fn backward(tensor: Self::FloatTensorPrimitive) -> Self::Gradients;
    fn grad(
        tensor: &Self::FloatTensorPrimitive,
        grads: &Self::Gradients,
    ) -> Option<<Self::InnerBackend as BackendTypes>::FloatTensorPrimitive>;
    fn grad_remove(
        tensor: &Self::FloatTensorPrimitive,
        grads: &mut Self::Gradients,
    ) -> Option<<Self::InnerBackend as BackendTypes>::FloatTensorPrimitive>;
    fn grad_replace(
        tensor: &Self::FloatTensorPrimitive,
        grads: &mut Self::Gradients,
        grad: <Self::InnerBackend as BackendTypes>::FloatTensorPrimitive,
    );
    fn inner(
        tensor: Self::FloatTensorPrimitive,
    ) -> <Self::InnerBackend as BackendTypes>::FloatTensorPrimitive;
    fn int_inner(
        tensor: Self::IntTensorPrimitive,
    ) -> <Self::InnerBackend as BackendTypes>::IntTensorPrimitive;
    fn bool_inner(
        tensor: Self::BoolTensorPrimitive,
    ) -> <Self::InnerBackend as BackendTypes>::BoolTensorPrimitive;
    fn q_inner(
        tensor: Self::QuantizedTensorPrimitive,
    ) -> <Self::InnerBackend as BackendTypes>::QuantizedTensorPrimitive;
    fn from_inner(
        tensor: <Self::InnerBackend as BackendTypes>::FloatTensorPrimitive,
    ) -> Self::FloatTensorPrimitive;
    fn int_from_inner(
        tensor: <Self::InnerBackend as BackendTypes>::IntTensorPrimitive,
    ) -> Self::IntTensorPrimitive;
    fn bool_from_inner(
        tensor: <Self::InnerBackend as BackendTypes>::BoolTensorPrimitive,
    ) -> Self::BoolTensorPrimitive;
    fn q_from_inner(
        tensor: <Self::InnerBackend as BackendTypes>::QuantizedTensorPrimitive,
    ) -> Self::QuantizedTensorPrimitive;

    // Provided methods
    fn set_distributed_params(
        tensor: Self::FloatTensorPrimitive,
        _param_id: ParamId,
    ) -> Self::FloatTensorPrimitive { ... }
    fn distributed_params(
        _tensor: &Self::FloatTensorPrimitive,
    ) -> Option<DistributedParams> { ... }
    fn is_distributed(_tensor: &Self::FloatTensorPrimitive) -> bool { ... }
}

Trait that allows a backend to support autodiff.

Required Associated Types

type InnerBackend: Backend<Device = Self::Device>

The inner backend type.

type Gradients: Send

Gradients type.

Required Methods

fn backward(tensor: Self::FloatTensorPrimitive) -> Self::Gradients

Backward pass.

Arguments

• tensor - The tensor is the last node of computational graph where the gradients are computed.

Returns

The gradients.

fn grad( tensor: &Self::FloatTensorPrimitive, grads: &Self::Gradients, ) -> Option<<Self::InnerBackend as BackendTypes>::FloatTensorPrimitive>

Returns the gradients of a tensor.

Arguments

• tensor - The tensor to extract the gradients from.

Returns

An optional tensor containing the gradient.

fn grad_remove( tensor: &Self::FloatTensorPrimitive, grads: &mut Self::Gradients, ) -> Option<<Self::InnerBackend as BackendTypes>::FloatTensorPrimitive>

Pops the gradients of a tensor and returns them.

Arguments

• tensor - The tensor to pop the gradients from.
• grads - The gradients.

Returns

An optional tensor containing the given gradients.

fn grad_replace( tensor: &Self::FloatTensorPrimitive, grads: &mut Self::Gradients, grad: <Self::InnerBackend as BackendTypes>::FloatTensorPrimitive, )

Replace the gradients of a tensor with the one provided.

If no gradient existed for the provided tensor, register it.

Arguments

• tensor - The tensor to pop the gradients from.
• grads - The gradients.
• grad - The updated grad tensor.

fn inner( tensor: Self::FloatTensorPrimitive, ) -> <Self::InnerBackend as BackendTypes>::FloatTensorPrimitive

Returns the tensor with inner backend type.

Arguments

• tensor - The tensor to get the inner backend tensor for.

Returns

The inner backend tensor.

fn int_inner( tensor: Self::IntTensorPrimitive, ) -> <Self::InnerBackend as BackendTypes>::IntTensorPrimitive

Returns the tensor with inner backend type.

Arguments

• tensor - The tensor to get the inner backend tensor for.

Returns

The inner backend tensor.

fn bool_inner( tensor: Self::BoolTensorPrimitive, ) -> <Self::InnerBackend as BackendTypes>::BoolTensorPrimitive

Returns the tensor with inner backend type.

Arguments

• tensor - The tensor to get the inner backend tensor for.

Returns

The inner backend tensor.

fn q_inner( tensor: Self::QuantizedTensorPrimitive, ) -> <Self::InnerBackend as BackendTypes>::QuantizedTensorPrimitive

Returns the tensor with inner backend type.

Arguments

• tensor - The tensor to get the inner backend tensor for.

Returns

The inner backend tensor.

fn from_inner( tensor: <Self::InnerBackend as BackendTypes>::FloatTensorPrimitive, ) -> Self::FloatTensorPrimitive

Converts the inner backend tensor to the autodiff backend tensor.

Arguments

• tensor - The inner backend tensor to convert.

Returns

The autodiff backend tensor.

fn int_from_inner( tensor: <Self::InnerBackend as BackendTypes>::IntTensorPrimitive, ) -> Self::IntTensorPrimitive

Converts the inner backend tensor to the autodiff backend tensor.

Arguments

• tensor - The inner backend tensor to convert.

Returns

The autodiff backend tensor.

fn bool_from_inner( tensor: <Self::InnerBackend as BackendTypes>::BoolTensorPrimitive, ) -> Self::BoolTensorPrimitive

Converts the inner backend tensor to the autodiff backend tensor.

Arguments

• tensor - The inner backend tensor to convert.

Returns

The autodiff backend tensor.

fn q_from_inner( tensor: <Self::InnerBackend as BackendTypes>::QuantizedTensorPrimitive, ) -> Self::QuantizedTensorPrimitive

Converts the inner backend tensor to the autodiff backend tensor.

Arguments

• tensor - The inner backend tensor to convert.

Returns

The autodiff backend tensor.

Provided Methods

fn set_distributed_params( tensor: Self::FloatTensorPrimitive, _param_id: ParamId, ) -> Self::FloatTensorPrimitive

Mark the tensor as distributed across multiple devices. The gradients will be aggregated during the backward pass.

This function does nothing when distributed training is not available.

fn distributed_params( _tensor: &Self::FloatTensorPrimitive, ) -> Option<DistributedParams>

Returns the distributed parameters if the tensor was marked as distributed.

fn is_distributed(_tensor: &Self::FloatTensorPrimitive) -> bool

Returns true if the tensor was marked as distributed.

Dyn Compatibility

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors

impl AutodiffBackend for Dispatch

Available on crate feature autodiff only.

type InnerBackend = Dispatch

type Gradients = Gradients

impl<B, C> AutodiffBackend for Autodiff<B, C>

where B: Backend, C: CheckpointStrategy,

type InnerBackend = B

type Gradients = Gradients