Trait rai_core::Func

pub trait Func<IN, OUT>
where
    IN: WithTensors,
    OUT: WithTensors,
{
    type Tangent: ToTensorGrads + FromTensorGrads;
    type Cotangent: ToTensorGrads + FromTensorGrads;

    // Required methods
    fn call(&self, input: IN) -> OUT;
    fn extract_input_tensors(&self, input: &IN, tensors: &mut Vec<Tensor>);

    // Provided method
    fn self_captured_tensors(&self, _tensors: &mut Vec<Tensor>) { ... }
}

Required Associated Types

type Tangent: ToTensorGrads + FromTensorGrads

type Cotangent: ToTensorGrads + FromTensorGrads

Required Methods

fn call(&self, input: IN) -> OUT

fn extract_input_tensors(&self, input: &IN, tensors: &mut Vec<Tensor>)

Provided Methods

fn self_captured_tensors(&self, _tensors: &mut Vec<Tensor>)

Implementors

impl<'m, 'i, 'l, M, F> Func<(&'m M, &'i Tensor, &'l Tensor), Tensor> for F

where M: Module, F: Fn(&'m M, &'i Tensor, &'l Tensor) -> Tensor,

type Tangent = HashMap<usize, Tensor>

type Cotangent = HashMap<usize, Tensor>

impl<'m, 'i, 'l, M, F, T> Func<(&'m M, &'i Tensor, &'l Tensor), (Tensor, Aux<T>)> for F

where M: Module, F: Fn(&'m M, &'i Tensor, &'l Tensor) -> (Tensor, Aux<T>),

type Tangent = HashMap<usize, Tensor>

type Cotangent = HashMap<usize, Tensor>

impl<'m, 'i, M, F> Func<(&'m M, &'i Tensor), Tensor> for F

where M: Module, F: Fn(&'m M, &'i Tensor) -> Tensor,

type Tangent = HashMap<usize, Tensor>

type Cotangent = HashMap<usize, Tensor>

impl<'m, 'i, M, F, T> Func<(&'m M, &'i Tensor), (Tensor, Aux<T>)> for F

where M: Module, F: Fn(&'m M, &'i Tensor) -> (Tensor, Aux<T>),

type Tangent = HashMap<usize, Tensor>

type Cotangent = HashMap<usize, Tensor>

impl<F> Func<[Tensor; 1], [Tensor; 1]> for F

where F: Fn(&Tensor) -> Tensor,

type Tangent = [Tensor; 1]

type Cotangent = [Tensor; 1]

impl<F> Func<[Tensor; 2], [Tensor; 1]> for F

where F: Fn(&Tensor, &Tensor) -> Tensor,

type Tangent = [Tensor; 2]

type Cotangent = [Tensor; 1]

impl<F> Func<[Tensor; 3], [Tensor; 1]> for F

where F: Fn(&Tensor, &Tensor, &Tensor) -> Tensor,

type Tangent = [Tensor; 3]

type Cotangent = [Tensor; 1]

impl<T> Func<Tensor, Tensor> for T

where T: Module,

type Tangent = HashMap<usize, Tensor>

type Cotangent = HashMap<usize, Tensor>