Function autograd::ops::argmax

pub fn argmax(x: &Tensor, axis: isize, keep_dim: bool) -> Tensor

Takes argmax along specified axis.

axis can be negative.

extern crate ndarray;
extern crate autograd as ag;

let mut ctx = ag::Context::new();
let ref x = ag::constant(ndarray::arr2(&[[3., 4.], [6., 5.]]), &mut ctx);
let ref y = ag::argmax(x, 1, false);

assert_eq!(y.eval(&mut ctx), ndarray::arr1(&[1., 0.]).into_dyn());