Implemented method 'create_array' for Base.ReshapedArrays.

JuliaSymbolics · Dec 11, 2024 · 4c16785 · 4c16785
1 parent 59e3aa6
commit 4c16785
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/src/code.jl b/src/code.jl
@@ -560,6 +560,10 @@ end
     create_array(P, S, nd, d, elems...)
 end
 
+@inline function create_array(::Type{<:Base.ReshapedArray{T, N, P}}, S, nd::Val, d, elems...) where {T, N, P}
+    create_array(P, S, nd, d, elems...)
+end
+
 ## SArray
 @inline function create_array(::Type{<:SArray}, ::Nothing, nd::Val, ::Val{dims}, elems...) where dims
     SArray{Tuple{dims...}}(elems...)

diff --git a/test/code.jl b/test/code.jl
@@ -150,6 +150,9 @@ nanmath_st.rewrites[:nanmath] = true
     @test eval(toexpr(Let([a ← 1, b ← 2, arr ← [1,2]],
                           MakeArray([a b;a+b a/b], arr)))) == [1 2; 3 1/2]
 
+    @test eval(toexpr(Let([a ← 1, b ← 2, arr ← [1,2]],
+                          MakeArray(reshape(view([a,b,a+b,a/b], :), 1, 4), arr)))) == [1 2 3 1/2]
+
     @test eval(toexpr(Let([a ← 1, b ← 2, arr ← @SVector([1,2])],
                           MakeArray([a,b,a+b,a/b], arr)))) === @SVector [1, 2, 3, 1/2]
 
@@ -159,6 +162,9 @@ nanmath_st.rewrites[:nanmath] = true
     @test eval(toexpr(Let([a ← 1, b ← 2, arr ← @SLVector((:a, :b))(@SVector[1,2])],
                           MakeArray([a+b,a/b], arr)))) === @SLVector((:a, :b))(@SVector [3, 1/2])
 
+    @test eval(toexpr(Let([a ← 1, b ← 2, arr ← reshape(view([1,2], :), 1, 2)],
+                          MakeArray([a,b,a+b,a/b], arr)))) == [1, 2, 3, 1/2]
+
     trackedarr = eval(toexpr(Let([a ← ReverseDiff.track(1.0), b ← 2, arr ← ReverseDiff.track(ones(2))],
                           MakeArray([a+b,a/b], arr))))
     @test trackedarr isa ReverseDiff.TrackedArray