1
1
use crate :: { shapes:: * , tensor:: * } ;
2
+ use std:: vec:: Vec ;
2
3
3
4
mod cpu_kernel;
4
5
#[ cfg( feature = "cuda" ) ]
@@ -21,15 +22,15 @@ mod webgpu_kernel;
21
22
/// # use dfdx_core::prelude::*;
22
23
/// # let dev: Cpu = Default::default();
23
24
/// let stack: Tensor<Rank3<2, 3, 4>, f32, _> = dev.zeros();
24
- /// let [a, b]: [Tensor<Rank2<3, 4>, f32, _>; 2] = stack.unstack();
25
+ /// let ( [a, b], _tape): ( [Tensor<Rank2<3, 4>, f32, _>; 2], _) = stack.unstack();
25
26
/// ```
26
27
///
27
28
/// Unstacking to a vec:
28
29
/// ```rust
29
30
/// # use dfdx_core::prelude::*;
30
31
/// # let dev: Cpu = Default::default();
31
- /// let stack: Tensor<(usize, Const::<3>, Const::<4>> , f32, _> = dev.zeros_like(&(2, Const, Const));
32
- /// let unstack: Vec<Tensor<Rank2<3, 4>, f32, _>> = stack.unstack();
32
+ /// let stack: Tensor<(usize, Const::<3>, Const::<4>) , f32, _> = dev.zeros_like(&(2, Const, Const));
33
+ /// let ( unstack, _tape): ( Vec<Tensor<Rank2<3, 4>, f32, _>>, _) = stack.unstack();
33
34
/// ```
34
35
pub trait TryUnstack < Head : Dim > : Sized {
35
36
type Unstacked ;
0 commit comments