Skip to content

Commit

Permalink
Only reconstruct with non-tuple T at the end for flattened output
Browse files Browse the repository at this point in the history
  • Loading branch information
serenity4 committed Nov 2, 2023
1 parent d2febca commit e607ea9
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 1 deletion.
8 changes: 7 additions & 1 deletion src/macro.jl
Original file line number Diff line number Diff line change
Expand Up @@ -554,7 +554,13 @@ function to_expr(cache, ex, flatten::Bool, T, variables, stop_early = false)
end
end
isexpr(ex, FACTOR) && return to_final_expr(cache, ex[1], flatten, T, variables)
isexpr(ex, KVECTOR) && return :($construct($(reconstructed_type(T, ex.cache.sig, ex)), $(Expr(:tuple, to_final_expr.(cache, ex, flatten, Ref(T), Ref(variables))...))))
if isexpr(ex, KVECTOR)
Tintermediate = flatten ? :Tuple : T
RT = reconstructed_type(Tintermediate, ex.cache.sig, ex)
components = Expr(:tuple)
append!(components.args, [to_final_expr(cache, x, flatten, Tintermediate, variables) for x in ex])
return :($construct($RT, $components))
end
isexpr(ex, BLADE) && return 1
if isexpr(ex, GEOMETRIC_PRODUCT)
@assert isweightedblade(ex)
Expand Down
8 changes: 8 additions & 0 deletions test/macro.jl
Original file line number Diff line number Diff line change
Expand Up @@ -212,4 +212,12 @@ using SymbolicGA: extract_weights, input_expression, extract_expression, restruc
@test (@ga 2 a::1 b::1 c::1 d::1) isa KVector
@test (@ga 2 a::1 (a::1 b::1) c::1 d::1) isa KVector
end

@testset "Flattening" begin
a, b, c, d = rand(3), rand(3), rand(3), rand(3)
ex = @ga 3 (a::1 b::1)::(0 + 2)
@test isa(ex, NTuple{4,Float64})
ex = @ga 3 NTuple{4,Float32} (a::1 b::1)::(0 + 2)
@test isa(ex, NTuple{4,Float32})
end
end;

0 comments on commit e607ea9

Please sign in to comment.