pub(crate) fn compute_gradients<'t, 'g, A, B, F: Float>(
    ys: &[A],
    wrt: &[B],
    gys: Option<&[Tensor<'g, F>]>, // not generic for None arg
    g: &'g Graph<F>,
) -> GradientMap<'g, F>
where
    A: AsRef<Tensor<'g, F>>,
    B: AsRef<Tensor<'g, F>>,
{
    // Setup gradient path.
    // We lookup this with tensor id.
    let mut between_nodes = get_between_nodes(g, ys, wrt);

    // Set default grads.
    if let Some(gys) = gys {
        assert_eq!(ys.len(), gys.len(), "`ys.len()` must match `gys.len()`");
        for (y, gy) in ys.into_iter().zip(gys) {
            between_nodes.get_mut(&y.as_ref().id).unwrap().default_grad = Some(gy.as_ref().clone());
        }
    } else {
        let start_gy = T::scalar(F::one(), g);
        for y in ys.into_iter() {
            between_nodes.get_mut(&y.as_ref().id).unwrap().default_grad = Some(start_gy.clone());
        }
    }

    // Prepare a heap with given ys.
    let mut heap = ys
        .iter()
        .map(|y| y.as_ref().to_node())
        .collect::<BinaryHeap<Node>>();

    // Backprop.
    // Starts with `ys`.
    while let Some(y) = heap.pop() {
        let gxs = {
            let info = between_nodes.get_mut(&y.id).unwrap();

            let gy = info.get_grad();

            // Call Op::grad (mutate the graph)
            let y_tensor = g.tensor(y.id);
            let gxs = GradientContext::new(gy, y_tensor, g).get_input_grads();
            debug_assert_eq!(y_tensor.num_backprop_inputs(), gxs.len());
            gxs
        };

        // Register computed gradients
        let y = g.tensor(y.id);
        for i in 0..gxs.len() {
            let x = y.get_backprop_input(i);
            let mut x_info = between_nodes.get_mut(&x.id).unwrap();
            if x_info.has_gradient {
                if let Some(gx) = gxs[i] {
                    x_info.push_grad(gx);
                    // update heap
                    if !x.is_source() && !x_info.grad_called {
                        x_info.grad_called = true;
                        heap.push(x.to_node());
                    }
                }
            }
        }
    }

    GradientMap {
        inner: between_nodes,
    }
}
/// Returns gradient tensors of `xs`.
///
/// This computes partial derivatives of `ys` with `xs` and returns the
/// gradients. This is achieved by building a subgraph between `ys` and
/// `xs` in reverse order from user's graph definition.
/// `gys` are already known gradients of `ys`'s outputs.
///
/// NOTE:
/// Returned gradient is `None` if the corresponding variable is not differentiable.
Returns gradient tensors of `xs`.
This computes partial derivatives of `ys` with `xs` and returns the
gradients. This is achieved by building a subgraph between `ys` and
`xs` in reverse order from user's graph definition.
`gys` are already known gradients of `ys`'s outputs.

Returned gradient is `None` if the corresponding variable is not differentiable.
Returns gradient tensors of ` xs ` . This computes partial derivatives of ` ys ` with ` xs ` and returns the gradients . This is achieved by building a subgraph between ` ys ` and ` xs ` in reverse order from user ' s graph definition . ` gys ` are already known gradients of ` ys ` ' s outputs . Returned gradient is ` None ` if the corresponding variable is not differentiable .
pub(crate) fn compute_gradients<'t, 'g, A, B, F: Float>(
    ys: &[A],
    wrt: &[B],
    gys: Option<&[Tensor<'g, F>]>, // not generic for None arg
    g: &'g Graph<F>,
) -> GradientMap<'g, F>
where
    A: AsRef<Tensor<'g, F>>,
    B: AsRef<Tensor<'g, F>>,
{
    // Setup gradient path.
    // We lookup this with tensor id.
    let mut between_nodes = get_between_nodes(g, ys, wrt);
    // Set default grads.
    if let Some(gys) = gys {
        assert_eq!(ys.len(), gys.len(), "`ys.len()` must match `gys.len()`");
        for (y, gy) in ys.into_iter().zip(gys) {
            between_nodes.get_mut(&y.as_ref().id).unwrap().default_grad = Some(gy.as_ref().clone());
        }
    } else {
        let start_gy = T::scalar(F::one(), g);
        for y in ys.into_iter() {
            between_nodes.get_mut(&y.as_ref().id).unwrap().default_grad = Some(start_gy.clone());
        }
    }
    // Prepare a heap with given ys.
    let mut heap = ys
        .iter()
        .map(|y| y.as_ref().to_node())
        .collect::<BinaryHeap<Node>>();
    // Backprop.
    // Starts with `ys`.
    while let Some(y) = heap.pop() {
        let gxs = {
            let info = between_nodes.get_mut(&y.id).unwrap();
            let gy = info.get_grad();
            // Call Op::grad (mutate the graph)
            let y_tensor = g.tensor(y.id);
            let gxs = GradientContext::new(gy, y_tensor, g).get_input_grads();
            debug_assert_eq!(y_tensor.num_backprop_inputs(), gxs.len());
            gxs
        };
        // Register computed gradients
        let y = g.tensor(y.id);
        for i in 0..gxs.len() {
            let x = y.get_backprop_input(i);
            let mut x_info = between_nodes.get_mut(&x.id).unwrap();
            if x_info.has_gradient {
                if let Some(gx) = gxs[i] {
                    x_info.push_grad(gx);
                    // update heap
                    if !x.is_source() && !x_info.grad_called {
                        x_info.grad_called = true;
                        heap.push(x.to_node());
                    }
                }
            }
        }
    }
    GradientMap {
        inner: between_nodes,
    }
}
pub ( crate ) fn compute_gradients < ' t , ' g , A , B , F : Float > ( ys : & [ A ] , wrt : & [ B ] , gys : Option < & [ Tensor < ' g , F > ] > , // not generic for None arg g : & ' g Graph < F > , ) -> GradientMap < ' g , F > where A : AsRef < Tensor < ' g , F > > , B : AsRef < Tensor < ' g , F > > , { // Setup gradient path. // We lookup this with tensor id. let mut between_nodes = get_between_nodes ( g , ys , wrt ) ; // Set default grads. if let Some ( gys ) = gys { assert_eq ! ( ys len ( ) gys len ( ) " " ) ; for ( y , gy ) in ys . into_iter ( ) . zip ( gys ) { between_nodes . get_mut ( & y . as_ref ( ) . id ) . unwrap ( ) . default_grad = Some ( gy . as_ref ( ) . clone ( ) ) ; } } else { let start_gy = T :: scalar ( F :: one ( ) , g ) ; for y in ys . into_iter ( ) { between_nodes . get_mut ( & y . as_ref ( ) . id ) . unwrap ( ) . default_grad = Some ( start_gy . clone ( ) ) ; } } // Prepare a heap with given ys. let mut heap = ys . iter ( ) . map ( | y | y . as_ref ( ) . to_node ( ) ) . collect :: < BinaryHeap < Node > > ( ) ; // Backprop. // Starts with `ys`. while let Some ( y ) = heap . pop ( ) { let gxs = { let info = between_nodes . get_mut ( & y . id ) . unwrap ( ) ; let gy = info . get_grad ( ) ; // Call Op::grad (mutate the graph) let y_tensor = g . tensor ( y . id ) ; let gxs = GradientContext :: new ( gy , y_tensor , g ) . get_input_grads ( ) ; debug_assert_eq ! ( y_tensor num_backprop_inputs ( ) gxs len ( ) ) ; gxs } ; // Register computed gradients let y = g . tensor ( y . id ) ; for i in 0 .. gxs . len ( ) { let x = y . get_backprop_input ( i ) ; let mut x_info = between_nodes . get_mut ( & x . id ) . unwrap ( ) ; if x_info . has_gradient { if let Some ( gx ) = gxs [ i ] { x_info . push_grad ( gx ) ; // update heap if ! x . is_source ( ) && ! x_info . grad_called { x_info . grad_called = true ; heap . push ( x . to_node ( ) ) ; } } } } } GradientMap { inner : between_nodes , } }
Returns gradient tensors of `xs`.
Returns gradient tensors of ` xs ` .
ys
&[A]
wrt
&[B]
gys
Option<&[Tensor<'g, F>]>
g
&'g Graph<F>
ys
&[A]
wrt
&[B]
gys
Option<&[Tensor<'g, F>]>
g
&'g Graph<F>