feat: optimize the special case of indexing with unitranges

EnzymeAD · Nov 9, 2024 · 8b90aea · 8b90aea
1 parent d786d2a
commit 8b90aea
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/ext/ReactantNNlibExt.jl b/ext/ReactantNNlibExt.jl
@@ -296,6 +296,17 @@ function NNlib.make_causal_mask(x::AnyTracedRArray; dims::Int=2)
     )
 end
 
+# XXX: reevaluate this manual optimization once
+#      https://github.com/EnzymeAD/Enzyme-JAX/issues/164 is handled
+function NNlib.gather!(
+    dst::TracedRArray{T1,2},
+    src::AnyTracedRArray{T2,2},
+    idxs::Union{AbstractUnitRange{<:Number}},
+) where {T1,T2}
+    dst.mlir_data = src[:, idxs].mlir_data
+    return dst
+end
+
 function NNlib.gather!(
     dst::TracedRArray{T1,2}, src::AnyTracedRArray{T2,2}, idxs::AbstractVector{<:Number}
 ) where {T1,T2}