fix: implement symbolic_type for ArrayOp and CallWithMetadata

JuliaSymbolics · Sep 5, 2024 · 299dba7 · 299dba7
1 parent fb0265e
commit 299dba7
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/src/arrays.jl b/src/arrays.jl
@@ -72,6 +72,8 @@ SymbolicUtils.sorted_arguments(s::ArrayOp) = sorted_arguments(s.term)
 
 shape(aop::ArrayOp) = aop.shape
 
+SymbolicIndexingInterface.symbolic_type(::Type{<:Symbolics.ArrayOp}) = ArraySymbolic()
+
 const show_arrayop = Ref{Bool}(false)
 function Base.show(io::IO, aop::ArrayOp)
     if iscall(aop.term) && !show_arrayop[]

diff --git a/src/variable.jl b/src/variable.jl
@@ -242,6 +242,8 @@ SymbolicUtils.Code.toexpr(x::CallWithMetadata, st) = SymbolicUtils.Code.toexpr(x
 
 CallWithMetadata(f) = CallWithMetadata(f, nothing)
 
+SymbolicIndexingInterface.symbolic_type(::Type{<:CallWithMetadata}) = ScalarSymbolic()
+
 function Base.show(io::IO, c::CallWithMetadata)
     show(io, c.f)
     print(io, "⋆")

diff --git a/test/symbolic_indexing_interface_trait.jl b/test/symbolic_indexing_interface_trait.jl
@@ -2,14 +2,18 @@ using Symbolics
 using SymbolicUtils
 using SymbolicIndexingInterface
 
-@test all(symbolic_type.([SymbolicUtils.BasicSymbolic, Symbolics.Num]) .==
+@test all(symbolic_type.([SymbolicUtils.BasicSymbolic, Symbolics.Num, Symbolics.CallWithMetadata]) .==
           (ScalarSymbolic(),))
-@test symbolic_type(Symbolics.Arr) == ArraySymbolic()
+@test all(symbolic_type.([Symbolics.ArrayOp, Symbolics.Arr]) .==
+          (ArraySymbolic(),))
 @variables x
 @test symbolic_type(x) == ScalarSymbolic()
 @variables y[1:3]
 @test symbolic_type(y) == ArraySymbolic()
 @test all(symbolic_type.(collect(y)) .== (ScalarSymbolic(),))
+@test symbolic_type(Symbolics.unwrap(y .* y)) == ArraySymbolic()
+@variables z(..)
+@test symbolic_type(z) == ScalarSymbolic()
 
 @variables x y z
 subs = Dict(x => 0.1, y => 2z)