! Construct a tensor with vec { 0 1 2 ... } and reshape to the desired shape
: naturals ( shape -- tensor )
- check-shape [ ] [ product [0..b) >float-array ] bi <tensor> ;
+ check-shape dup product [0..b) >float-array <tensor> ;
! Construct a tensor without initializing its values
: (tensor) ( shape -- tensor )
! Flatten the tensor so that it is only one-dimensional
: flatten ( tensor -- tensor )
- dup shape>>
- product { } 1sequence >>shape ;
+ dup shape>> product { } 1sequence >>shape ;
! outputs the number of dimensions of a tensor
: dims ( tensor -- n )
! Sum together all elements in the tensor
syntax:M: tensor sum vec>> 0 <simd-slice>
- [ simd-slice>> 0 [ sum + ] reduce ]
+ [ simd-slice>> [ sum ] map-sum ]
[ end-slice>> sum ] bi + ;
<PRIVATE
:: final-vstack-shape ( seq -- shape )
! Compute the new second-to-last dimension
seq first dims 2 - :> vdim
- seq 0 [ shape>> vdim swap nth + ] reduce
+ seq [ shape>> vdim swap nth ] map-sum
! Combine it to create the new shape
seq first shape>> clone :> new-shape
vdim new-shape set-nth
! Compute the final shape
[
! Compute the first dimension
- [ 0 [ shape>> first + ] reduce 1array ]
+ [ [ shape>> first ] map-sum 1array ]
! Compute the other dimensions
[ first shape>> rest ] bi append
]