Function autograd::ops::tensordot [] [src]

pub fn tensordot<T: ArrayLike>(
    a: &Tensor,
    b: &Tensor,
    a_axes: &T,
    b_axes: &T
) -> Tensor

Computes tensor dot product (tensor contraction) along specified axes.

Arguments

  • a - Input tensor
  • b - Input tensor
  • a_axes - Contraction axes
  • b_axes - Contraction axes

Note1: length of a_axes and b_axes must match.

Note2: Each axis number can be negative.

extern crate autograd as ag;

let mut ctx = ag::Context::new();

let ref a = ag::zeros(&[3, 4, 5]);
let ref b = ag::zeros(&[4, 3, 2]);
let ref c = ag::tensordot(a, b, &[1, 0], &[0, 1]);
assert_eq!(c.eval(&mut ctx).shape(), &[5, 2]);

For detailed description, see https://docs.scipy.org/doc/numpy/reference/generated/numpy.tensordot.html.