Add tests for extended operators

Author: MilesCranmer

test/test_error_handling.jl

@@ -7,11 +7,12 @@ T = Union{baseT,Vector{baseT},Matrix{baseT}}
 
 scalar_add(x::T, y::T) where {T<:Real} = x + y
 
-operators = GenericOperatorEnum(; binary_operators=[scalar_add], extend_user_operators=true)
+operators = GenericOperatorEnum(; binary_operators=[scalar_add])
 
 x1, x2, x3 = [Node(T; feature=i) for i in 1:3]
 
-tree = Node(1, x1, x2)
+@extend_operators operators
+tree = scalar_add(x1, x2)
 
 # With error handling:
 try

test/test_tensor_operators.jl

@@ -8,7 +8,7 @@ function vec_add(x, y)
     return x .+ y
 end
 
-operators = GenericOperatorEnum(; binary_operators=[vec_add], extend_user_operators=true)
+operators = GenericOperatorEnum(; binary_operators=[vec_add])
 
 x1, x2, x3 = [Node(T; feature=i) for i in 1:3]
 c1 = Node(T; val=[1.0, 2.0, 3.0])
@@ -22,20 +22,30 @@ 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]
 
+# Try same things, but with constructors:
+@extend_operators operators
+tree = vec_add(c1, x2)
+@test repr(tree) == "vec_add([1.0, 2.0, 3.0], x2)"
+@test tree(X) == [4.0, 5.0, 6.0]
+tree = vec_add(x1, c1)
+@test repr(tree) == "vec_add(x1, [1.0, 2.0, 3.0])"
+@test tree(X) == [3.0, 4.0, 5.0]
+
 # Also test unary operators:
 function vec_square(x)
     return x .* x
 end
 
-operators = GenericOperatorEnum(;
-    binary_operators=[vec_add], unary_operators=[vec_square], extend_user_operators=true
-)
+operators = GenericOperatorEnum(; binary_operators=[vec_add], unary_operators=[vec_square])
+@extend_operators operators
 tree = Node(1, c1)
 @test repr(tree) == "vec_square([1.0, 2.0, 3.0])"
 @test tree(X) == [1.0, 4.0, 9.0]
+@test vec_square(c1).val == [1.0, 4.0, 9.0]
 tree = Node(1, Node(1, c1), x1)
 @test repr(tree) == "vec_add(vec_square([1.0, 2.0, 3.0]), x1)"
 @test tree(X) == [3.0, 6.0, 11.0]
+@test (vec_add(vec_square(c1), x1))(X) == [3.0, 6.0, 11.0]
 
 # Also test mixed scalar and floats:
 c2 = Node(T; val=2.0)

test/test_tree_construction.jl

@@ -106,3 +106,20 @@ s = String(take!(io))
 set_node!(tree, tree2)
 @test tree !== tree2
 @test repr(tree) == repr(tree2)
+
+# Test that we can work with custom operators:
+function op1(x, y)
+    return x + y
+end
+function op2(x, y)
+    return x ^ 2 + 1/((y)^2 + 0.1)
+end
+function op3(x)
+    return sin(x) + cos(x)
+end
+operators = OperatorEnum(; default_params..., binary_operators=(op1, op2), unary_operators=(op3,))
+@extend_operators operators
+x1 = Node(; feature=1)
+x2 = Node(; feature=2)
+tree = op1(op2(x1, x2), op3(x1))
+@test repr(tree) == "op1(op2(x1, x2), op3(x1))"