extra/tensors/tensors.factor

   1 ! Copyright (C) 2019 HMC Clinic.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3
   4 USING: accessors alien alien.c-types alien.data arrays combinators
   5 grouping kernel math math.functions ranges math.vectors
   6 math.vectors.simd multi-methods parser prettyprint.custom sequences sequences.extras
   7 sequences.private specialized-arrays typed ;
   8
   9 QUALIFIED-WITH: alien.c-types c
  10 SPECIALIZED-ARRAY: c:float
  11 SPECIALIZED-ARRAY: float-4
  12 IN: tensors
  13
  14 ! Tensor class definition
  15 TUPLE: tensor
  16     { shape array }
  17     { vec float-array } ;
  18
  19 ! Errors
  20 ERROR: non-positive-shape-error shape ;
  21 ERROR: shape-mismatch-error shape1 shape2 ;
  22 ERROR: non-uniform-seq-error seq ;
  23 ERROR: dimension-mismatch-error tensor-dim index-dim ;
  24
  25 <PRIVATE
  26
  27 ! Check that the shape has only positive values
  28 : check-shape ( shape -- shape )
  29     dup [ 1 < ] map-find drop [ non-positive-shape-error ] when ;
  30
  31 ! Construct a tensor of zeros
  32 : <tensor> ( shape seq -- tensor )
  33     tensor boa ;
  34
  35 ! Creates a freshly-allocated float-array with the desired c-type values
  36 : >float-array ( seq -- float-array )
  37     c:float >c-array ;
  38
  39 : repetition ( shape const -- tensor )
  40     [ check-shape dup product ] dip <repetition>
  41     >float-array <tensor> ;
  42
  43 PRIVATE>
  44
  45 ! Construct a tensor of zeros
  46 : zeros ( shape -- tensor )
  47     0 repetition ;
  48
  49 ! Construct a tensor of ones
  50 : ones ( shape -- tensor )
  51     1 repetition ;
  52
  53 ! Construct a one-dimensional tensor with values start, start+step,
  54 ! ..., stop (inclusive)
  55 : arange ( a b step -- tensor )
  56     <range> [ length >fixnum 1array ] keep >float-array <tensor> ;
  57
  58 ! Construct a tensor with vec { 0 1 2 ... } and reshape to the desired shape
  59 : naturals ( shape -- tensor )
  60     check-shape [ ] [ product [0..b) >float-array ] bi <tensor> ;
  61
  62 ! Construct a tensor without initializing its values
  63 : (tensor) ( shape -- tensor )
  64     dup product (float-array) <tensor> ;
  65
  66 <PRIVATE
  67
  68 : check-reshape ( shape1 shape2 -- shape1 shape2 )
  69     2dup [ product ] bi@ = [ shape-mismatch-error ] unless ;
  70
  71 PRIVATE>
  72
  73 ! Reshape the tensor to conform to the new shape
  74 : reshape ( tensor shape -- tensor )
  75     [ dup shape>> ] [ check-shape ] bi* check-reshape nip >>shape ;
  76
  77 ! Flatten the tensor so that it is only one-dimensional
  78 : flatten ( tensor -- tensor )
  79     dup shape>>
  80     product { } 1sequence >>shape ;
  81
  82 ! outputs the number of dimensions of a tensor
  83 : dims ( tensor -- n )
  84     shape>> length ;
  85
  86 ! Turn into Factor ND array form
  87 ! Source: shaped-array>array
  88 TYPED: tensor>array ( tensor: tensor -- seq: array )
  89     [ vec>> >array ] [ shape>> ] bi
  90     [ rest-slice reverse [ group ] each ] unless-empty ;
  91
  92 <PRIVATE
  93 ! recursively finds shape of nested array
  94 ! assumes properly shaped array (all sub-arrays are same size)
  95 :: find-shape ( seq shape -- shape' )
  96     seq empty? [ { 0 } ] [
  97         ! add length of seq element to shape
  98         shape seq length 1array append :> shape'
  99         ! base case: check if the first element is a seq
 100         seq first :> 1st
 101         1st sequence?
 102         ! is a sequence: recurse on 1st element
 103         [ 1st shape' find-shape ]
 104         ! not a sequence: return shape'
 105         [ shape' ] if
 106     ] if ;
 107 PRIVATE>
 108
 109 ! turns a nested array into a tensor
 110 :: >tensor ( seq -- tensor )
 111     ! get the shape
 112     seq { } find-shape :> shape
 113     ! flatten the array
 114     seq
 115     shape length 1 - [
 116         drop concat
 117     ] each-integer :> flatseq
 118     ! check that the size is good
 119     shape product flatseq length =
 120     [ seq non-uniform-seq-error ] unless
 121     ! turn into a tensor
 122     shape flatseq >float-array <tensor> ;
 123
 124 SYNTAX: t{ \ } [ >tensor ] parse-literal ;
 125
 126 ! Pretty printing
 127 syntax:M: tensor pprint-delims drop \ t{ \ } ;
 128 syntax:M: tensor >pprint-sequence tensor>array ;
 129 syntax:M: tensor pprint* pprint-object ;
 130
 131
 132 <PRIVATE
 133 ! turns a shape into a list of things by which to multiply
 134 ! indices to get a full index (e.g. { 2 3 4 } -> { 12 4 1 })
 135 : ind-mults ( shape -- seq )
 136     <reversed> 1 swap [ swap [ * ] keep ] map nip reverse ;
 137
 138 ! turns a num/seq index & tensor into num index & tensor
 139 ! also throws a dimension mismatch if seq & tens shape>> arent the same len
 140 : num-index ( n/seq tensor -- n tensor )
 141     ! check form of index (num or seq)
 142     swap dup array? not
 143     [ ! if array, first check if it's a valid index
 144         2dup [ shape>> length ] dip length 2dup =
 145         [ dimension-mismatch-error ] unless 2drop
 146         ! turn into num
 147         [ dup shape>> ind-mults ] dip [ * ] 2map-sum
 148     ] unless swap ;
 149
 150 PRIVATE>
 151
 152
 153 ! Sequence protocol implementation
 154 syntax:M: tensor clone [ shape>> clone ] [ vec>> clone ] bi <tensor> ;
 155
 156 syntax:M: tensor length vec>> length ;
 157
 158 syntax:M: tensor nth num-index vec>> nth ;
 159
 160 syntax:M: tensor nth-unsafe num-index vec>> nth-unsafe ;
 161
 162 syntax:M: tensor set-nth num-index vec>> set-nth ;
 163
 164 syntax:M: tensor set-nth-unsafe num-index vec>> set-nth-unsafe ;
 165
 166 syntax:M: tensor new-sequence
 167     ! Check if the old and new tensors are the same size
 168     shape>> 2dup product =
 169     ! If so preserve the shape, otherwise create a 1D tensor
 170     [ nip (tensor) ] [ drop 1array (tensor) ] if ;
 171
 172 syntax:M: tensor like
 173     ! If the original sequence is already a tensor, we are done
 174     over tensor?
 175     [ drop ] [
 176         over float-array? [
 177             [ dup [ length 1array ] dip <tensor> ] dip
 178         ] [
 179             [ >tensor ] dip
 180         ] if
 181         2dup [ length ] bi@ = [ shape>> reshape ] [ drop ] if
 182     ] if ;
 183
 184 syntax:M: tensor clone-like
 185     ! If the original sequence is already a tensor, we just need to clone it
 186     over tensor?
 187     [ drop clone ] [
 188         [ >tensor ] dip
 189         2dup [ length ] bi@ = [ shape>> reshape ] [ drop ] if
 190     ] if ;
 191
 192 INSTANCE: tensor sequence
 193
 194
 195 <PRIVATE
 196
 197 :: make-subseq ( arr start len -- arr )
 198     ! Find the index
 199     c:float heap-size start *
 200     ! Compute the starting pointer
 201     arr underlying>> <displaced-alien>
 202     ! Push length and type to create the new array
 203     len c:float <c-direct-array> ; inline
 204
 205 : check-bop-shape ( shape1 shape2 -- shape )
 206     2dup = [ shape-mismatch-error ] unless drop ;
 207
 208 ! Apply the binary operator bop to combine the tensors
 209 TYPED:: t-bop ( tensor1: tensor tensor2: tensor quot: ( x y -- z ) -- tensor: tensor )
 210     tensor1 shape>> tensor2 shape>> check-bop-shape
 211     tensor1 vec>> tensor2 vec>> quot 2map <tensor> ; inline
 212
 213 ! Create an array of 4-element SIMD arrays for processing floats
 214 : simd-for-bop ( array -- simd-array rest-slice/f )
 215     dup length dup 4 mod [ drop f ] [ - cut-slice ] if-zero
 216     [ float-4 cast-array ] dip ; inline
 217
 218 ! Create an array of 4-element SIMD arrays for processing floats
 219 ! Tensor class definition
 220 TUPLE: simd-slice
 221     { first-slice float-array }
 222     { simd-slice float-4-array }
 223     { end-slice float-array } ;
 224
 225 :: (simd-slice) ( arr start len -- arr/f )
 226     len [ float-array{ } ] [ drop arr start len make-subseq ] if-zero ; inline
 227
 228 :: <simd-slice> ( arr start -- simd-slice )
 229     ! Compute the beginning
 230     arr 0 start (simd-slice)
 231     ! Compute the SIMD part
 232     arr length start - :> len
 233     len 4 mod :> end
 234     arr start len end - (simd-slice) float-4 cast-array
 235     ! Compute the end
 236     arr dup length end - end (simd-slice)
 237     simd-slice boa ; inline
 238
 239 ! Apply the binary operators simd-quot and quot to quickly combine the tensors
 240 :: t-bop-simd ( tensor1 tensor2 simd-quot: ( x y -- z ) quot: ( x y -- z ) -- tensor )
 241     tensor1 shape>> tensor2 shape>> check-bop-shape
 242     tensor1 vec>> tensor2 vec>>
 243     dup length (float-array) dup :> vec3
 244     [ simd-for-bop ] tri@ :> ( simd1 rest1 simd2 rest2 simd3 rest3 )
 245     simd1 simd2 simd-quot simd3 2map-into
 246     rest1 rest2 quot rest3 2map-into
 247     vec3 <tensor> ; inline
 248
 249 ! Apply the operation to the tensor
 250 TYPED:: t-uop ( tensor: tensor quot: ( x -- y ) -- tensor: tensor )
 251     tensor vec>> quot map [ tensor shape>> ] dip <tensor> ; inline
 252
 253 ! Apply the binary operators simd-quot and quot to quickly combine a tensor and
 254 ! a number
 255 :: t-uop-simd ( tensor n simd-quot: ( x y -- z ) quot: ( x y -- z ) -- tensor )
 256     tensor dup [ shape>> ] [ vec>> ] bi*
 257     dup length (float-array) dup :> vec2
 258     [ simd-for-bop ] bi@ :> ( simd1 rest1 simd2 rest2 )
 259     simd1 n n n n float-4-boa simd-quot curry simd2 map-into
 260     rest1 n quot curry rest2 map-into
 261     vec2 <tensor> ; inline
 262
 263 PRIVATE>
 264
 265 ! Add a tensor to either another tensor or a scalar
 266 multi-methods:GENERIC: t+ ( x y -- tensor )
 267 METHOD: t+ { tensor tensor } [ v+ ] [ + ] t-bop-simd ;
 268 METHOD: t+ { tensor number } >float [ v+ ] [ + ] t-uop-simd ;
 269 METHOD: t+ { number tensor } swap >float [ swap v+ ] [ swap + ] t-uop-simd ;
 270
 271 ! Subtraction between two tensors or a tensor and a scalar
 272 multi-methods:GENERIC: t- ( x y -- tensor )
 273 METHOD: t- { tensor tensor } [ v- ] [ - ] t-bop-simd ;
 274 METHOD: t- { tensor number } >float [ v- ] [ - ] t-uop-simd ;
 275 METHOD: t- { number tensor } swap >float [ swap v- ] [ swap - ] t-uop-simd ;
 276
 277 ! Multiply a tensor with either another tensor or a scalar
 278 multi-methods:GENERIC: t* ( x y -- tensor )
 279 METHOD: t* { tensor tensor } [ v* ] [ * ] t-bop-simd ;
 280 METHOD: t* { tensor number } >float [ v* ] [ * ] t-uop-simd ;
 281 METHOD: t* { number tensor } swap >float [ swap v* ] [ swap * ] t-uop-simd ;
 282
 283 ! Divide two tensors or a tensor and a scalar
 284 multi-methods:GENERIC: t/ ( x y -- tensor )
 285 METHOD: t/ { tensor tensor } [ v/ ] [ / ] t-bop-simd ;
 286 METHOD: t/ { tensor number } >float [ v/ ] [ / ] t-uop-simd ;
 287 METHOD: t/ { number tensor } swap >float [ swap v/ ] [ swap / ] t-uop-simd ;
 288
 289 ! Mod two tensors or a tensor and a scalar
 290 multi-methods:GENERIC: t% ( x y -- tensor )
 291 METHOD: t% { tensor tensor } [ mod ] t-bop ;
 292 METHOD: t% { tensor number } >float [ mod ] curry t-uop ;
 293 METHOD: t% { number tensor } [ >float ] dip [ mod ] with t-uop ;
 294
 295 ! Sum together all elements in the tensor
 296 syntax:M: tensor sum vec>> 0 <simd-slice>
 297     [ simd-slice>> 0 [ sum + ] reduce ]
 298     [ end-slice>> sum ] bi + ;
 299
 300 <PRIVATE
 301
 302 ! Also converts all elements of the sequence to tensors
 303 :: check-concat-shape ( seq -- seq )
 304     ! Compute the bottom shape of the first element in the sequence
 305     seq first { } >tensor dup :> empty-tensor
 306     like shape>> dup :> first-shape rest :> rest-shape
 307     seq [
 308         ! Compute the bottom shape of this element
 309         empty-tensor like dup shape>> rest
 310         ! Compare; if they are different, throw an error
 311         rest-shape = [ shape>> first-shape swap shape-mismatch-error ] unless
 312     ] map ;
 313
 314 ! Also converts all elements of the sequence to tensors
 315 :: check-stack-shape ( seq -- seq )
 316     ! Compute the bottom shape of the first element in the sequence
 317     seq first { } >tensor dup :> empty-tensor
 318     like shape>> :> first-shape
 319     seq [
 320         ! Compute the bottom shape of this element
 321         empty-tensor like dup shape>>
 322         ! Compare; if they are different, throw an error
 323         first-shape = [ shape>> first-shape swap shape-mismatch-error ] unless
 324     ] map ;
 325
 326 ! Also converts all elements of the sequence to tensors
 327 :: check-hstack-shape ( seq -- seq )
 328     ! Compute the top shape of the first element in the sequence
 329     seq first { } >tensor dup :> empty-tensor
 330     like shape>> dup :> first-shape but-last :> but-last-shape
 331     seq [
 332         ! Compute the top shape of this element
 333         empty-tensor like dup shape>> but-last
 334         ! Compare; if they are different, throw an error
 335         but-last-shape = [ shape>> first-shape swap shape-mismatch-error ] unless
 336     ] map ;
 337
 338 : final-hstack-shape ( seq -- shape )
 339     ! Get the top part
 340     dup first shape>> but-last swap
 341     ! Compute the last part of the shape
 342     [ shape>> last ] map sum 1array append ;
 343
 344 ! Returns an guide for hstacking where the index corresponds to the postion
 345 ! in the last dimension of the resulting tensor, and the elements are
 346 ! { which tensor, len of tensor, index }
 347 :: hstack-guide ( seq -- guide )
 348     ! Compute the list of last shape parts
 349     seq [ shape>> last ] map :> last-dims
 350     ! Curr tensor and index in tensor
 351     0 0
 352     last-dims sum [0..b) [
 353         drop :> old-t-ind :> last-dims-i
 354         last-dims-i last-dims nth
 355         old-t-ind -
 356         ! If we need to move onto the next tensor
 357         [ last-dims-i 1 + 0 ]
 358         ! Otherwise, stay with the current tensor
 359         [ drop last-dims-i old-t-ind ] if-zero
 360         2dup [ dup last-dims nth ] dip 3array
 361         [ 1 + ] dip
 362     ] map nip nip ;
 363
 364 ! Given a sequence of tensors, stack them across the last dimension
 365 :: hstack-unsafe ( tseq -- tensor )
 366     ! Create the final tensor
 367     tseq final-hstack-shape (tensor)
 368     ! Compute the guide information
 369     tseq hstack-guide dup length :> repeat :> guide
 370     dup vec>> [
 371         :> i drop
 372         ! First get the correct tensor
 373         i repeat /mod guide nth
 374         dup first tseq nth
 375         ! Now find the correct value within that tensor
 376         [ [ second ] [ third ] bi -rot * + ] dip nth
 377     ] map-index! drop ;
 378
 379 ! Also converts all elements of the sequence to tensors
 380 :: check-vstack-shape ( seq -- seq )
 381     ! Compute the shape of the first sequence
 382     seq first { } >tensor dup :> empty-tensor
 383     like shape>> dup :> first-shape
 384     ! Compute the index of the dimension to be stacked across
 385     length 2 - :> vdim
 386     seq [
 387         ! Convert this element to a tensor
 388         empty-tensor like dup
 389         ! Compare the shapes
 390         shape>> first-shape [ = ] 2map
 391         vdim swap remove-nth
 392         ! If the shapes differ in anything except the second-to-last dimension
 393         ! this sequence cannot be vstacked
 394         t [ = ] reduce [ shape>> first-shape swap shape-mismatch-error ] unless
 395     ] map ;
 396
 397 ! Compute the shape after the vstack has been completed
 398 :: final-vstack-shape ( seq -- shape )
 399     ! Compute the new second-to-last dimension
 400     seq first dims 2 - :> vdim
 401     seq 0 [ shape>> vdim swap nth + ] reduce
 402     ! Combine it to create the new shape
 403     seq first shape>> clone :> new-shape
 404     vdim new-shape set-nth
 405     new-shape ;
 406
 407 ! Combine the second-to-last and last dimensions of each tensor for stacking
 408 :: reshape-for-vstack ( seq -- seq )
 409     seq first dims 2 - :> vdim
 410     seq [
 411         dup shape>> vdim cut product 1array append >>shape
 412     ] map! ;
 413
 414
 415 PRIVATE>
 416
 417 ! Concatenation operations
 418 ! Concatenate across the last dimension
 419 : t-concat ( seq -- tensor )
 420     check-concat-shape
 421     ! Compute the final shape
 422     [
 423         ! Compute the first dimension
 424         [ 0 [ shape>> first + ] reduce 1array ]
 425         ! Compute the other dimensions
 426         [ first shape>> rest ] bi  append
 427     ]
 428     ! Concatenate all of the float-arrays
 429     [ [ vec>> ] map concat ] bi <tensor> ;
 430
 431 : stack ( seq -- tensor )
 432     check-stack-shape
 433     ! Compute the new shape
 434     [ [ length 1array ] [ first shape>> ] bi append ]
 435     ! Concatenate all of the tensors
 436     [ [ vec>> ] map concat ] bi <tensor> ;
 437
 438 : hstack ( seq -- tensor )
 439     ! Check shape and convert everything to tensors
 440     check-hstack-shape hstack-unsafe ;
 441
 442 : vstack ( seq -- tensor )
 443     ! Check shape and convert everything to tensors
 444     check-vstack-shape
 445     ! Find the final shape
 446     [ final-vstack-shape ]
 447     ! Reshape each of the tensors and stack
 448     [ reshape-for-vstack hstack-unsafe ] bi
 449     ! Finally reshape and return
 450     swap >>shape ;
 451
 452 <PRIVATE
 453
 454 ! Check that the tensor has an acceptable shape for matrix multiplication
 455 : check-matmul-shape ( tensor1 tensor2 -- )
 456     [let [ shape>> ] bi@ :> shape2 :> shape1
 457     ! Check that the matrices can be multiplied
 458     shape1 last shape2 [ length 2 - ] keep nth =
 459     ! Check that the other dimensions are equal
 460     shape1 2 head* shape2 2 head* = and
 461     ! If either is false, raise an error
 462     [ shape1 shape2 shape-mismatch-error ] unless ] ;
 463
 464 ! Slice out a row from the array
 465 : row ( arr n i p -- slice )
 466     ! Compute the starting index
 467     / truncate dupd *
 468     ! Compute the ending index
 469     swap over +
 470     ! Take a slice
 471     rot <slice> ;
 472
 473 ! much quicker transpose for 2d tensors
 474 TYPED:: 2d-transpose ( tensor: tensor -- tensor': tensor )
 475     tensor shape>> :> old-shape
 476     tensor vec>> :> vec
 477     old-shape first2 :> ( s1 s2 )
 478     ! loop through new tensor
 479     old-shape reverse dup product <iota> [
 480         ! find y*b val in original tensor
 481         s1 /mod s2 *
 482         ! find x val in original tensor
 483         [ s2 /mod ] dip + nip
 484         ! get that index in original tensor
 485         vec nth-unsafe
 486     ] float-array{ } map-as <tensor> ;
 487
 488 ! Perform matrix multiplication muliplying an
 489 ! mxn matrix with a nxp matrix
 490 TYPED:: 2d-matmul ( vec1: float-array vec2: float-array res: float-array
 491                     m: fixnum n: fixnum p: fixnum -- )
 492     ! For each element in the range, we want to compute the dot product of the
 493     ! corresponding row and column
 494     ! Transpose vec2 so that we are doing row * row (as opposed to row * col)
 495     { n p } vec2 <tensor> 2d-transpose vec>> :> vec2
 496
 497     m [ :> i
 498         i n * :> in
 499         i p * :> ip
 500         vec1 in n make-subseq
 501         p [ :> j
 502             dup
 503             vec2 j n * n make-subseq
 504             0.0 [ * + ] 2reduce
 505             ip j + res set-nth-unsafe
 506         ] each-integer
 507         drop
 508     ] each-integer ;
 509
 510 ! Perform matrix multiplication muliplying an
 511 ! mxn matrix with a nxp matrix
 512 TYPED:: 2d-matmul-mixed ( vec1: float-array vec2: float-array res: float-array
 513                     m: fixnum n: fixnum p: fixnum start: fixnum -- )
 514     ! For each element in the range, we want to compute the dot product of the
 515     ! corresponding row and column
 516     ! Transpose vec2 so that we are doing row * row (as opposed to row * col)
 517     { n p } vec2 <tensor> 2d-transpose vec>> :> vec2
 518
 519     ! Compute the location in the float-array each 2D matrix will start at
 520     start m n * * :> start1
 521     start n p * * :> start2
 522
 523     m [ :> i
 524         i n * :> in
 525         4 4 in start1 + 4 mod - swap mod :> in4m
 526         i p * :> ip
 527         vec1 in n make-subseq :> sub1
 528         sub1 in4m <simd-slice> :> slice1
 529         p [ :> j
 530             j n * :> jn
 531             4 4 jn 4 mod - swap mod :> jn4m
 532             vec2 jn n make-subseq
 533             in4m jn4m = [
 534                 jn4m <simd-slice> slice1 swap
 535                 2dup [ first-slice>> ] bi@ 0.0 [ * + ] 2reduce
 536                 [ 2dup [ simd-slice>> ] bi@ ] dip [ vdot + ] 2reduce
 537                 [ [ end-slice>> ] bi@ ] dip [ * + ] 2reduce
 538             ] [
 539                 sub1 swap
 540                 0.0 [ * + ] 2reduce
 541             ] if
 542             ip j + res set-nth-unsafe
 543         ] each-integer
 544     ] each-integer ;
 545
 546 ! ! Perform matrix multiplication muliplying an
 547 ! mxn matrix with a nxp matrix
 548 ! Should only be called when n is a multiple of 4
 549 TYPED:: 2d-matmul-simd ( vec1: float-array vec2: float-array
 550                              res: float-array
 551                              m: fixnum n: fixnum p: fixnum -- )
 552     ! For each element in the range, we want to compute the dot product of the
 553     ! corresponding row and column
 554     ! Transpose vec2 so that we are doing row * row (as opposed to row * col)
 555     { n p } vec2 <tensor> 2d-transpose vec>> :> vec2
 556
 557     m [ :> i
 558         i n * :> in
 559         i p * :> ip
 560         vec1 in n make-subseq float-4 cast-array
 561         p [ :> j
 562             dup
 563             vec2 j n * n make-subseq float-4 cast-array
 564             0.0 [ vdot + ] 2reduce
 565             ip j + res set-nth-unsafe
 566         ] each-integer
 567         drop
 568     ] each-integer ;
 569
 570 PRIVATE>
 571
 572
 573 ! Perform matrix multiplication muliplying an
 574 ! ...xmxn matrix with a ...xnxp matrix
 575 TYPED:: matmul ( tensor1: tensor tensor2: tensor -- tensor3: tensor )
 576     ! First check the shape
 577     tensor1 tensor2 check-matmul-shape
 578
 579     ! Now save all of the sizes
 580     tensor1 shape>> unclip-last-slice :> n
 581     unclip-last-slice :> m :> top-shape
 582     tensor2 shape>> last :> p
 583     top-shape product :> top-prod
 584
 585     ! Create the shape of the resulting tensor
 586     top-shape { m p } append
 587
 588     ! Now create the new float array to store the underlying result
 589     dup product (float-array) :> vec3
 590
 591     ! Now update the tensor3 to contain the multiplied matricies
 592     top-prod [
 593         :> i
 594
 595         ! Compute vec1 using direct C arrays
 596         tensor1 vec>> m n * i * m n * make-subseq
 597
 598         ! Compute vec2 and start2
 599         tensor2 vec>> n p * i * n p * make-subseq
 600
 601         ! Compute the result
 602         vec3 m p * i * m p * make-subseq
 603         ! Push m, n, and p and multiply the arrays
 604         m n p
 605         { { [ n 4 mod 0 = ] [ 2d-matmul-simd ] }
 606           { [ n 4 < ] [ 2d-matmul ] }
 607           [ i 2d-matmul-mixed ]
 608         } cond
 609
 610     ] each-integer
 611     vec3 <tensor> ;
 612
 613 ! Transpose an n-dimensional tensor by flipping the axes
 614 TYPED:: transpose ( tensor: tensor -- tensor': tensor )
 615     tensor shape>> length 2 =
 616     [ tensor 2d-transpose ]
 617     [ tensor shape>> :> old-shape
 618         tensor vec>> :> vec
 619         old-shape reverse :> new-shape
 620         old-shape ind-mults :> mults
 621         ! loop through new tensor
 622         new-shape dup product <iota> [
 623             ! find index in original tensor
 624             old-shape mults [ [ /mod ] dip * ] 2map-sum nip
 625             ! get that index in original tensor
 626             vec nth-unsafe
 627         ] float-array{ } map-as <tensor>
 628     ] if ;