Include tensor operator test

Author: MilesCranmer

src/Equation.jl

@@ -83,7 +83,12 @@ function Base.convert(
     get!(id_map, tree) do
         if tree.degree == 0
             if tree.constant
-                Node(0, tree.constant, convert(T1, (tree.val::T2)))
+                val = tree.val::T2
+                if !(T2 <: T1)
+                    # e.g., we don't want to convert Float32 to Union{Float32,Vector{Float32}}!
+                    val = convert(T1, val)
+                end
+                Node(T1, 0, tree.constant, val)
             else
                 Node(T1, 0, tree.constant, nothing, tree.feature)
             end
@@ -138,6 +143,10 @@ function Node(
             "You must specify either `val` or `feature` when creating a leaf node, not both.",
         )
     elseif T2 <: Nothing
+        if !(T1 <: T)
+            # Only convert if not already in the type union.
+            val = convert(T, val)
+        end
         return Node(T, 0, true, val)
     else
         return Node(T, 0, false, nothing, feature)

test/test_tensor_operators.jl

@@ -0,0 +1,23 @@
+using DynamicExpressions
+using Test
+
+baseT = Float64
+T = Union{baseT,Vector{baseT},Matrix{baseT}}
+
+function vec_add(x, y)
+    return x .+ y
+end
+
+operators = GenericOperatorEnum(; binary_operators=[vec_add], extend_user_operators=true)
+
+x1, x2, x3 = [Node(T; feature=i) for i in 1:3]
+c1 = Node(T; val=[1.0, 2.0, 3.0])
+
+X = [[2.0, 2.0, 2.0], [3.0, 3.0, 3.0]]
+
+tree = Node(1, c1, x2)
+@test repr(tree) == "vec_add([1.0, 2.0, 3.0], x2)"
+@test tree(X) == [4.0, 5.0, 6.0]
+tree = Node(1, x1, c1)
+@test repr(tree) == "vec_add(x1, [1.0, 2.0, 3.0])"
+@test tree(X) == [3.0, 4.0, 5.0]

test/unittest.jl

@@ -47,3 +47,7 @@ end
 @safetestset "Test generic operators" begin
     include("test_generic_operators.jl")
 end
+
+@safetestset "Test tensor operators" begin
+    include("test_tensor_operators.jl")
+end