Function dfdx::tensor_ops::max_axis

source · [−]
pub fn max_axis<T: Reduce1<I>, const I: isize>(t: T) -> T::Reduced
Expand description

Reduces dimension I of the tensor by gathering the maximum value from that dimension.

Pytorch equivalent: t.amax(I)

NOTE This evenly distributes gradients between all equal maximum values, instead of only exactly 1 value.

Examples:

let t = Tensor2D::new([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
let r: Tensor1D<2> = t.max_axis::<-1>();
assert_eq!(r.data(), &[3.0, -1.0]);