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]);