feat: overload circshift (#1386)

* feat: overload circshift

* test: fftshift

Author: avik-pal

src/TracedRArray.jl

@@ -1410,4 +1410,27 @@ function Base._reverse!(a::AnyTracedRArray{T,N}, dims::NTuple{M,Int}) where {T,N
     return a
 end
 
+function Base.circshift!(
+    dest::AnyTracedRArray{T,N}, src, shiftamt::Base.DimsInteger
+) where {T,N}
+    src = TracedUtils.promote_to(TracedRArray{T,N}, materialize_traced_array(src))
+    shiftamt = Base.fill_to_length(shiftamt, 0, Val(N))
+
+    for i in 1:N
+        amt = shiftamt[i] % size(src, i)
+        amt == 0 && continue
+        if amt > 0
+            src1 = selectdim(src, i, (size(src, i) - amt + 1):size(src, i))
+            src2 = selectdim(src, i, 1:(size(src, i) - amt))
+        else
+            src1 = selectdim(src, i, (-amt + 1):size(src, i))
+            src2 = selectdim(src, i, 1:(-amt))
+        end
+        src = cat(src1, src2; dims=i)
+    end
+
+    copyto!(dest, src)
+    return dest
+end
+
 end

test/basic.jl

@@ -1337,3 +1337,13 @@ sameunitrange(x, y) = first(x) == first(y) && last(x) == last(y)
         end
     end
 end
+
+@testset "circshift" begin
+    x = reshape(collect(Float32, 1:36), 2, 6, 3)
+    x_ra = Reactant.to_rarray(x)
+
+    @test @jit(circshift(x_ra, (1, 2))) ≈ circshift(x, (1, 2))
+    @test @jit(circshift(x_ra, (1, 2, 3))) ≈ circshift(x, (1, 2, 3))
+    @test @jit(circshift(x_ra, (-3, 2))) ≈ circshift(x, (-3, 2))
+    @test @jit(circshift(x_ra, (5, 2))) ≈ circshift(x, (5, 2))
+end

test/integration/fft.jl

@@ -21,6 +21,10 @@ using FFTW, Reactant, Test
     y_ra = @jit(fft(x_ra))
     @test @jit(ifft(y_ra)) ≈ x
 
+    shifted_fft = @jit(fftshift(y_ra))
+    @test shifted_fft ≈ fftshift(Array(y_ra))
+    @test @jit(ifftshift(shifted_fft)) ≈ Array(y_ra)
+
     @testset "fft real input" begin
         x = rand(Float32, 2, 3, 4)
         x_ra = Reactant.to_rarray(x)