toki/numidr
1
0
Fork 0
Code Issues Pull requests Activity

Add more utility functions

Browse source
This commit is contained in:
Kiana Sheibani 2022-11-30 22:02:56 -05:00
parent 524486bb58
commit eddb0e318c
Signed by: toki
GPG key ID: 6CB106C25E86A9F7
3 changed files with 33 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

19
src/Data/NumIdr/Array/Array.idr
View file

@ -516,6 +516,25 @@ stack axis arrs = rewrite sym (lengthCorrect arrs) in
getAxisInd FZ (i :: is) = (i, is)
getAxisInd {s=_::_} (FS ax) (i :: is) = mapSnd (i::) (getAxisInd ax is)
export
joinAxes : {s' : _} -> Array s (Array s' a) -> Array (s ++ s') a
joinAxes {s} arr with (viewShape arr)
_ | Shape s = fromFunctionNB (s ++ s') (\is => arr !# takeUpTo s is !# dropUpTo s is)
where
takeUpTo : Vect rk Nat -> Vect (rk + rk') Nat -> Vect rk Nat
takeUpTo [] ys = []
takeUpTo (x::xs) (y::ys) = y :: takeUpTo xs ys
dropUpTo : Vect rk Nat -> Vect (rk + rk') Nat -> Vect rk' Nat
dropUpTo [] ys = ys
dropUpTo (x::xs) (y::ys) = dropUpTo xs ys
export
splitAxes : (rk : Nat) -> {0 rk' : Nat} -> {s : _} ->
Array {rk=rk+rk'} s a -> Array (take {m=rk'} rk s) (Array (drop {m=rk'} rk s) a)
splitAxes _ {s} arr = fromFunctionNB _ (\is => fromFunctionNB _ (\is' => arr !# (is ++ is')))
||| Construct the transpose of an array by reversing the order of its axes.
export