Function autograd::ops::split [] [src]

pub fn split(x: &Tensor, sizes: &[usize], axis: isize) -> Vec<Tensor>

Splits input tensors into parts.

Splits x into sizes.len() parts along axis.

The size of dimension of each part is sizes[i] on axis, but x.shape[i] on other axis.

extern crate autograd as ag;

let ref a = ag::zeros(&[3, 7, 5]);
let ref b = ag::split(a, &[2, 3, 2], 1);

let evaluated = ag::eval(&[&b[0], &b[1], &b[2]], &mut ag::Context::new());

assert_eq!(evaluated[0].shape(), &[3, 2, 5]);
assert_eq!(evaluated[1].shape(), &[3, 3, 5]);
assert_eq!(evaluated[2].shape(), &[3, 2, 5]);