feat: support batched transpose and adjoint

Author: avik-pal

ext/ReactantNNlibExt.jl

@@ -203,4 +203,7 @@ function NNlib.meanpool(x::AnyTracedRArray{T}, pdims::NNlib.PoolDims) where {T}
            T(numel)
 end
 
+NNlib.batched_transpose(x::AnyTracedRArray{T,3}) where {T} = permutedims(x, (2, 1, 3))
+NNlib.batched_adjoint(x::AnyTracedRArray{<:Real,3}) = NNlib.batched_transpose(x)
+
 end # module ReactantNNlibExt

test/wrapped_arrays.jl

@@ -1,4 +1,4 @@
-using Reactant, Test, Statistics
+using Reactant, Test, Statistics, NNlib
 
 function view_getindex_1(x)
     x = view(x, 2:3, 1:2, :)
@@ -98,3 +98,18 @@ end
     @test m1 ≈ m2
     @test v1 ≈ v2
 end
+
+function btranspose_badjoint(x)
+    x1 = NNlib.batched_transpose(x)
+    x2 = NNlib.batched_adjoint(x)
+    return x1 .+ x2
+end
+
+@testset "batched transpose/adjoint" begin
+    x = rand(4, 2, 3)
+    x_ra = Reactant.to_rarray(x)
+
+    btranspose_badjoint_compiled = @compile btranspose_badjoint(x_ra)
+
+    @test btranspose_badjoint_compiled(x_ra) ≈ btranspose_badjoint(x)
+end