feat: broadcasting on views

Author: avik-pal

src/TracedRArray.jl

@@ -680,6 +680,7 @@ function Base.mapreducedim!(
 end
 
 struct AbstractReactantArrayStyle{N} <: Base.Broadcast.AbstractArrayStyle{N} end
+
 AbstractReactantArrayStyle(::Val{N}) where {N} = AbstractReactantArrayStyle{N}()
 AbstractReactantArrayStyle{M}(::Val{N}) where {N,M} = AbstractReactantArrayStyle{N}()
 
@@ -690,7 +691,9 @@ AbstractReactantArrayStyle{M}(::Val{N}) where {N,M} = AbstractReactantArrayStyle
 #    copy(inst)
 # end
 
-BroadcastStyle(::Type{T}) where {T<:TracedRArray} = AbstractReactantArrayStyle{ndims(T)}()
+function BroadcastStyle(::Type{<:AnyTracedRArray{T,N}}) where {T,N}
+    return AbstractReactantArrayStyle{N}()
+end
 
 function Base.similar(
     bc::Broadcasted{AbstractReactantArrayStyle{N}}, ::Type{T}, dims
@@ -758,6 +761,12 @@ function broadcast_to_size(arg::AbstractArray, rsize)
     return arg
 end
 
+function broadcast_to_size(
+    arg::WrappedArray{T,N,TracedRArray,<:TracedRArray{T,N}}, rsize
+) where {T,N}
+    return broadcast_to_size(arg[axes(arg)...], rsize)
+end
+
 function broadcast_to_size(arg::TracedRArray, rsize)
     return arg
 end

test/wrapped_arrays.jl

@@ -48,3 +48,35 @@ end
 
     @test reshape_wrapper_compiled(x_ra) ≈ reshape_wrapper(x)
 end
+
+function permutedims_wrapper(x)
+    x = view(x, 2:3, 1:2, :)
+    return permutedims(x, (2, 1, 3))
+end
+
+@testset "permutedims wrapper" begin
+    x = rand(4, 4, 3)
+    x_ra = Reactant.to_rarray(x)
+
+    permutedims_wrapper(x)
+
+    permutedims_wrapper_compiled = @compile permutedims_wrapper(x_ra)
+
+    @test permutedims_wrapper_compiled(x_ra) ≈ permutedims_wrapper(x)
+end
+
+function bcast_wrapper(f::F, x) where {F}
+    x = view(x, 2:3, :)
+    return f.(x)
+end
+
+@testset "Broadcasting on wrapped arrays" begin
+    x = rand(4, 3)
+    x_ra = Reactant.to_rarray(x)
+
+    for op in (-, tanh, sin)
+        bcast_compiled = @compile bcast_wrapper(op, x_ra)
+
+        @test bcast_compiled(op, x_ra) ≈ bcast_wrapper(op, x)
+    end
+end