feat: handle mixed types in body of template expression

Author: MilesCranmer

src/ComposableExpression.jl

@@ -260,22 +260,45 @@ end
 # Basically we want to vectorize every single operation on ValidVector,
 # so that the user can use it easily.
 
+function _apply_operator(op::F, x::Vararg{Any,N}) where {F<:Function,N}
+    vx = map(_get_value, x)
+    safe_op = get_safe_op(op)
+    result = safe_op.(vx...)
+    return ValidVector(result, is_valid_array(result))
+end
+
 function apply_operator(op::F, x::Vararg{Any,N}) where {F<:Function,N}
     if all(_is_valid, x)
-        vx = map(_get_value, x)
-        safe_op = get_safe_op(op)
-        result = safe_op.(vx...)
-        return ValidVector(result, is_valid_array(result))
+        return _apply_operator(op, x...)
     else
         example_vector =
             something(map(xi -> xi isa ValidVector ? xi : nothing, x)...)::ValidVector
-        return ValidVector(_get_value(example_vector), false)
+        expected_return_type = Base.promote_op(
+            _apply_operator, typeof(op), map(typeof, x)...
+        )
+        if expected_return_type !== Union{} &&
+            expected_return_type <: ValidVector{<:AbstractArray}
+            return ValidVector(
+                _match_eltype(expected_return_type, example_vector.x), false
+            )::expected_return_type
+        else
+            return ValidVector(example_vector.x, false)
+        end
     end
 end
 _is_valid(x::ValidVector) = x.valid
 _is_valid(x) = true
 _get_value(x::ValidVector) = x.x
 _get_value(x) = x
+function _match_eltype(
+    ::Type{<:ValidVector{<:AbstractArray{T1}}}, x::AbstractArray{T2}
+) where {T1,T2}
+    if T1 == T2
+        return x
+    else
+        return Base.Fix1(convert, T1).(x)
+    end
+end
 
 struct ValidVectorMixError <: Exception end
 struct ValidVectorAccessError <: Exception end

src/TemplateExpression.jl

@@ -665,6 +665,18 @@ and this would automatically handle the validity and vectorization.
     )
 end
 
+function _match_input_eltype(
+    ::Type{<:AbstractMatrix{T1}}, result::AbstractVector{T2}
+) where {T1,T2}
+    if T1 != T2 && T1 <: AbstractFloat && T2 <: AbstractFloat
+        # Just to handle cases where the user might write
+        # 0.5 in their template spec, but the data is Float32.
+        return Base.Fix1(convert, T1).(result)
+    else
+        return result
+    end
+end
+
 @stable(
     default_mode = "disable",
     default_union_limit = 2,
@@ -695,7 +707,7 @@ end
             if !(result isa ValidVector)
                 throw(TemplateReturnError())
             end
-            return result.x, result.valid
+            return _match_input_eltype(typeof(cX), result.x), result.valid
         end
         function (ex::TemplateExpression)(
             X, operators::Union{AbstractOperatorEnum,Nothing}=nothing; kws...

test/test_composable_expression.jl

@@ -182,3 +182,23 @@ end
     @test ex(ValidVector([1.0, 1.0], true), 2.0).x ≈ [3.0, 3.0]
     @test ex(ValidVector([1.0, 1.0], false), 2.0).valid == false
 end
+
+@testitem "ValidVector operations with Union{} return type" tags = [:part2] begin
+    using SymbolicRegression: ValidVector
+    using SymbolicRegression.ComposableExpressionModule: apply_operator
+
+    error_op(::Any, ::Any) = error("This should cause Union{} inference")
+
+    x = ValidVector([1.0, 2.0], false)
+    y = ValidVector([3.0, 4.0], false)
+
+    result = apply_operator(error_op, x, y)
+    @test result isa ValidVector
+    @test !result.valid
+    @test result.x == [1.0, 2.0]
+
+    a = ValidVector(Float32[1.0, 2.0], false)
+    b = 1.0
+    result2 = apply_operator(*, a, b)
+    @test result2 isa ValidVector{<:AbstractArray{Float64}}
+end

test/test_template_expression.jl

@@ -692,3 +692,31 @@ end
     @test contains(msg, "ValidVector is required")
     @test contains(msg, "ValidVector(my_data, computation_is_valid)")
 end
+
+@testitem "Test Float32/Float64 type conversion in TemplateExpression" tags = [:part2] begin
+    using SymbolicRegression
+    using SymbolicRegression: eval_loss
+
+    template = @template_spec(expressions = (f,)) do x1, x2
+        0.5 * f(x1, x2)  # 0.5 is Float64 literal
+    end
+
+    options = Options(; binary_operators=[+, *, /, -], expression_spec=template)
+    x1 = ComposableExpression(Node{Float32}(; feature=1); operators=options.operators)
+    x2 = ComposableExpression(Node{Float32}(; feature=2); operators=options.operators)
+    f_expr = x1 + x2
+
+    template_expr = TemplateExpression(
+        (; f=f_expr); structure=template.structure, operators=options.operators
+    )
+
+    X = Float32[1.0 2.0; 3.0 4.0]
+    result = template_expr(X)
+    @test result isa Vector{Float32}
+
+    y = Float32[2.0, 3.0]
+    dataset = Dataset(X, y)
+    loss = eval_loss(template_expr, dataset, options)
+    @test loss isa Float32
+    @test loss ≈ 0.0
+end