Update sugar apis (#1844)

* Update sugar apis * cleanup * cleanup * cleanup * fix * fix * fix stack * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * Update internal_rules.jl * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix * fix
EnzymeAD · Sep 17, 2024 · 786a998 · 786a998
1 parent bd5dcd1
commit 786a998
Show file tree

Hide file tree

Showing 5 changed files with 609 additions and 495 deletions.
diff --git a/ext/EnzymeStaticArraysExt.jl b/ext/EnzymeStaticArraysExt.jl
@@ -3,6 +3,8 @@ module EnzymeStaticArraysExt
 using StaticArrays
 using Enzyme
 
+@inline Enzyme.tupstack(rows::(NTuple{N, <:StaticArrays.SArray} where N), inshape, outshape) = reshape(cat(rows..., dims=length(inshape)), (inshape..., outshape...)) 
+
 @inline function Enzyme.onehot(x::StaticArrays.SArray{S, T, N, L}) where {S, T, N, L}
     ntuple(Val(L)) do i
         Base.@_inline_meta