Merge pull request #482 from MilesCranmer/auto-convert-template

feat: handle mixed types in body of template expression

Author: MilesCranmer

.github/workflows/CI.yml

@@ -86,7 +86,30 @@ jobs:
           julia --color=yes --threads=auto --check-bounds=yes --depwarn=yes --code-coverage=user -e 'import Coverage; import Pkg; Pkg.activate("."); Pkg.test(coverage=true)'
           julia --color=yes coverage.jl
         shell: bash
-      - uses: codecov/codecov-action@v5
+      - name: "Upload coverage artifacts"
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-${{ matrix.julia-version }}-${{ matrix.os }}-${{ matrix.test }}
+          path: lcov.info
+          retention-days: 1
+
+  upload-coverage:
+    name: Upload Coverage to Codecov
+    needs: test
+    runs-on: ubuntu-latest
+    # Only run on pushes to master or pull requests
+    if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name != github.repository
+    steps:
+      - uses: actions/checkout@v5
+      - name: Download all coverage artifacts
+        uses: actions/download-artifact@v4
+        with:
+          pattern: coverage-*
+          path: coverage
+      - name: Upload to Codecov
+        uses: codecov/codecov-action@v5
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
-          files: lcov.info
+          directory: ./coverage
+          fail_ci_if_error: true
+          verbose: true

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,36 @@ 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, _match_eltype
+
+    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}}
+
+    # Test apply_operator when all inputs are valid
+    valid_x = ValidVector([1.0, 2.0], true)
+    valid_y = ValidVector([3.0, 4.0], true)
+    valid_result = apply_operator(+, valid_x, valid_y)
+    @test valid_result.valid == true
+    @test valid_result.x ≈ [4.0, 6.0]
+
+    # cover _match_eltype
+    arr = [1.0, 2.0]
+    @test _match_eltype(ValidVector{Vector{Float64}}, arr) === arr  # Same type
+    arr_f32 = Float32[1.0, 2.0]
+    @test _match_eltype(ValidVector{Vector{Float64}}, arr_f32) isa Vector{Float64}  # Different type
+end

test/test_template_expression.jl

@@ -692,3 +692,38 @@ 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
+    using SymbolicRegression.TemplateExpressionModule: _match_input_eltype
+
+    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
+
+    # Test _match_input_eltype coverage (covers lines 675-676)
+    result_f64 = [1.0, 2.0]
+    @test _match_input_eltype(Matrix{Float64}, result_f64) === result_f64  # Same type
+    result_int = [1, 2]
+    @test _match_input_eltype(Matrix{Float64}, result_int) === result_int  # Non-float type
+end