apply_type_with_promotion #1681

jumerckx · 2025-09-18T08:36:38Z

While in my testing traced_type now does the typevar promotion we talked about, it doesn't yet work for make_tracer as the result from traced_types isn't used to guide make_tracer AFAICT.

using Reactant

struct Bar{T}
    b::T
end
struct Foo{T, B<:Bar{T}, AT<:AbstractArray{T}}
    a::AT
    b::B
end

foo = Reactant.to_rarray(Foo(rand(3), Bar(1.)))

Reactant.traced_type(typeof(foo), Val(Reactant.ConcreteToTraced), Union{}, nothing, nothing)
# returns: Foo{Reactant.TracedRNumber{Float64}, Bar{Reactant.TracedRNumber{Float64}}, Reactant.TracedRArray{Float64, 1}}

Reactant.make_tracer(Reactant.OrderedIdDict(), foo, (), Reactant.ConcreteToTraced) # fails

ERROR: AssertionError: Could not recursively make tracer of object of type Foo{Float64, Bar{Float64}, ConcretePJRTArray{Float64, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}} into Foo{Reactant.TracedRNumber{Float64}, Bar{Reactant.TracedRNumber{Float64}}, Reactant.TracedRArray{Float64, 1}} at field 2 (named b), need object of type Bar{Reactant.TracedRNumber{Float64}} found object of type Bar{Float64} 
Stacktrace:
 [1] make_tracer_unknown(seen::Reactant.OrderedIdDict{…}, prev::Any, path::Any, mode::Reactant.TraceMode; track_numbers::Type, sharding::Any, runtime::Any, kwargs::@Kwargs{})
   @ Reactant ~/Reactant2/src/Tracing.jl:1197
 [2] make_tracer_unknown
   @ ~/Reactant2/src/Tracing.jl:1095 [inlined]
 [3] #make_tracer#139
   @ ~/Reactant2/src/Tracing.jl:1232 [inlined]
 [4] make_tracer(seen::Reactant.OrderedIdDict{Any, Any}, prev::Any, path::Any, mode::Reactant.TraceMode)
   @ Reactant ~/Reactant2/src/Tracing.jl:1222
 [5] top-level scope
   @ ~/Reactant2/make_tracer_debuggin.jl:12
Some type information was truncated. Use `show(err)` to see complete types.

wsmoses · 2025-09-18T12:58:04Z

src/Tracing.jl

+                # The field is constrained by a TypeVar directly,
+                # so we don't need to check.
+                # (The check below would fail if the typevar was promoted as
+                # we don't get the same result when calling traced_type_inner


don't we still need to check below that the subtype matches, even if constrained? just to confirm we did do the typevar solve correctly?

I'm not sure...
The reason I added this early return is for e.g.:
Foo{Float64, Bar{Float64}, ConcretePJRTArray{Float64, 1}}
Which, with ConcreteToTraced is now converted to:
Foo{TracedRNumber{Float64}, Bar{TracedRNumber{Float64}}, TracedRArray{Float64, 1}}

The field b of the original type is Bar{Float64}. When tracing this, it just returns Bar{Float64} because of lack of constraints. So what should we compare the fieldtype against?

struct T d::Array end

struct T{D<:Array} d::D end struct T{D<:AbstractArray{Float64}} d::D end

check can be skipped when an unconstrained typevar

must be skipped when that typevar is itself a dependency of another type. If this is true in any case, we need the fancy make_tracer

may be run when that typevar is itself not a dependency of another type. If this is true in all cases, we don't need the upgrade

case where need to skip checj

struct MyNewArray{FT, N, D <: AbstractArray{FT, N}} data::D off::FT end

wsmoses · 2025-09-18T12:58:24Z

src/Tracing.jl

+This function tries to apply the param types to the wrapper type.
+When there's a constraint conflict, it tries to resolve it by promoting the conflicting types. The new param type is then propagated in any param type that depends on it.
+"""
+function apply_type_with_promotion(wrapper, params, relevant_typevars=typevar_dict(wrapper))


can we add some tests?

…ype promotion

github-actions · 2025-11-29T18:07:55Z

src/Tracing.jl

+function collect_tvars_in_type!(dependencies, @nospecialize(t))
+    if t isa TypeVar
+        push!(dependencies, t)
+        return


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

return

return nothing

github-actions · 2025-11-29T18:07:55Z

src/Tracing.jl

        return T
    end

+    @debug "traced_type_inner: Processing type with field changes" T=T subTys=subTys


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

@debug "traced_type_inner: Processing type with field changes" T=T subTys=subTys

@debug "traced_type_inner: Processing type with field changes" T = T subTys = subTys

github-actions · 2025-11-29T18:07:56Z

src/Tracing.jl

    wrapped_cifrt_array = T <: AbstractArray && ancestor(T) <: ConcreteIFRTArray
    wrapped_tracedarray = T <: AbstractArray && ancestor(T) <: TracedRArray

+    @debug "wrapped flags" wrapped_cpjrt_array=wrapped_cpjrt_array wrapped_cifrt_array=wrapped_cifrt_array wrapped_tracedarray=wrapped_tracedarray


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

@debug "wrapped flags" wrapped_cpjrt_array=wrapped_cpjrt_array wrapped_cifrt_array=wrapped_cifrt_array wrapped_tracedarray=wrapped_tracedarray

@debug "wrapped flags" wrapped_cpjrt_array = wrapped_cpjrt_array wrapped_cifrt_array =

wrapped_cifrt_array wrapped_tracedarray = wrapped_tracedarray

github-actions · 2025-11-29T18:07:56Z

src/Tracing.jl

+    @debug "wrapped flags" wrapped_cpjrt_array=wrapped_cpjrt_array wrapped_cifrt_array=wrapped_cifrt_array wrapped_tracedarray=wrapped_tracedarray
+
    subParms = []
+    @debug "Tracing type parameters" num_params=length(T.parameters) T_parameters=T.parameters


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

@debug "Tracing type parameters" num_params=length(T.parameters) T_parameters=T.parameters

@debug "Tracing type parameters" num_params = length(T.parameters) T_parameters =

T.parameters

github-actions · 2025-11-29T18:07:56Z

src/Tracing.jl

        end
    end

+    @debug "Built subParms" subParms=subParms


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

@debug "Built subParms" subParms=subParms

@debug "Built subParms" subParms = subParms

github-actions · 2025-11-29T18:07:56Z

src/Tracing.jl

            return TT2
        end
+    else
+        @debug "Field count mismatch" fieldcount_T=fieldcount(T) fieldcount_TT2=fieldcount(TT2)


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

@debug "Field count mismatch" fieldcount_T=fieldcount(T) fieldcount_TT2=fieldcount(TT2)

@debug "Field count mismatch" fieldcount_T = fieldcount(T) fieldcount_TT2 = fieldcount(

TT2

)

github-actions · 2025-11-29T18:07:56Z

src/Tracing.jl

                    xi2 = Core.Typeof(xi2)((newpath,), xi2.mlir_data)
                    seen[xi2] = xi2
                    changed = true
+                elseif !ismutabletype(FT) && !ismutabletype(Core.Typeof(xi2)) && fieldcount(FT) == fieldcount(Core.Typeof(xi2))


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

elseif !ismutabletype(FT) && !ismutabletype(Core.Typeof(xi2)) && fieldcount(FT) == fieldcount(Core.Typeof(xi2))

elseif !ismutabletype(FT) &&

!ismutabletype(Core.Typeof(xi2)) &&

fieldcount(FT) == fieldcount(Core.Typeof(xi2))

github-actions · 2025-11-29T18:07:57Z

src/Tracing.jl

+                            val_wrapped = ft_j(val_j)
+                            # Correct the path for the wrapped scalar
+                            sub_path = append_path(newpath, j)
+                            val_wrapped = Core.Typeof(val_wrapped)((sub_path,), val_wrapped.mlir_data)


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

val_wrapped = Core.Typeof(val_wrapped)((sub_path,), val_wrapped.mlir_data)

val_wrapped = Core.Typeof(val_wrapped)(

(sub_path,), val_wrapped.mlir_data

)

github-actions · 2025-11-29T18:07:57Z

src/Tracing.jl

+                    end
+
+                    if success
+                        xi2 = ccall(:jl_new_structv, Any, (Any, Ptr{Any}, UInt32), FT, flds_sub, fieldcount(FT))


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

xi2 = ccall(:jl_new_structv, Any, (Any, Ptr{Any}, UInt32), FT, flds_sub, fieldcount(FT))

xi2 = ccall(

:jl_new_structv,

Any,

(Any, Ptr{Any}, UInt32),

FT,

flds_sub,

fieldcount(FT),

)

test/tracing.jl

+            ) == (Foo{
+                TracedRNumber{Float64},
+                Bar{TracedRNumber{Float64}},
+                Reactant.TracedRArray{Float64,1},
+            }, [true, true, false])


jumerckx · 2025-11-30T21:37:11Z

I added a bunch of debug statements that seem handy to keep around.
However, they clutter CI. Should I just get rid of them, or look for a solution?

wsmoses · 2025-12-01T19:13:31Z

src/Tracing.jl

        return T
    end

+    @debug "traced_type_inner: Processing type with field changes" T = T subTys = subTys


I would probably get rid of the debug statements here, if no longer necessary

wsmoses · 2025-12-01T19:14:33Z

@giordano adding you as a reviewer.

this should remedy many of the custom type annotations weve historically needed for oceananigans and fixedsizearray

jumerckx · 2025-12-01T20:00:34Z

Not sure how I should fix the CI docs failure. It's an internal method but still useful to have a docstring imo.

giordano · 2025-12-02T00:28:45Z

Add it to https://enzymead.github.io/Reactant.jl/stable/api/api#Internal-utils

apply_type_with_promotion

b4e1455

jumerckx requested a review from wsmoses September 18, 2025 08:36

wsmoses reviewed Sep 18, 2025

View reviewed changes

jumerckx added 6 commits September 20, 2025 00:03

add test

12b1d63

test formatting

138b9db

remove new_params

c9b4f12

update make_tracer to try correct fieldtypes according to the fancy t…

0faa1dc

…ype promotion

Merge branch 'main' into jm/promote_type_params

f6e3226

bail out for non-type fields

e2141ef

jumerckx mentioned this pull request Nov 22, 2025

TypeError when compiling a function involving a struct #1885

Open

jumerckx added 3 commits November 29, 2025 12:03

wip

32f75f5

fix test

d213c55

get rid of try-catch

fca90a0

github-actions bot reviewed Nov 29, 2025

View reviewed changes

format

29a3534

jumerckx marked this pull request as ready for review November 29, 2025 18:21

jumerckx requested a review from wsmoses December 1, 2025 16:52

wsmoses reviewed Dec 1, 2025

View reviewed changes

wsmoses requested a review from giordano December 1, 2025 19:14

wsmoses requested a review from avik-pal December 1, 2025 19:14

jumerckx added 2 commits December 1, 2025 13:31

get rid of @debug statements

5866b84

example in docstring

65f9d24

add to docs

7edef34

	@debug "traced_type_inner: Processing type with field changes" T=T subTys=subTys
	@debug "traced_type_inner: Processing type with field changes" T = T subTys = subTys

	@debug "wrapped flags" wrapped_cpjrt_array=wrapped_cpjrt_array wrapped_cifrt_array=wrapped_cifrt_array wrapped_tracedarray=wrapped_tracedarray
	@debug "wrapped flags" wrapped_cpjrt_array = wrapped_cpjrt_array wrapped_cifrt_array =
	wrapped_cifrt_array wrapped_tracedarray = wrapped_tracedarray

	@debug "Tracing type parameters" num_params=length(T.parameters) T_parameters=T.parameters
	@debug "Tracing type parameters" num_params = length(T.parameters) T_parameters =
	T.parameters

	@debug "Built subParms" subParms=subParms
	@debug "Built subParms" subParms = subParms

apply_type_with_promotion #1681

Are you sure you want to change the base?

apply_type_with_promotion #1681

Conversation

jumerckx commented Sep 18, 2025

Uh oh!

wsmoses Sep 18, 2025

Choose a reason for hiding this comment

Uh oh!

jumerckx Sep 19, 2025

Choose a reason for hiding this comment

Uh oh!

wsmoses Sep 24, 2025

Choose a reason for hiding this comment

Uh oh!

wsmoses Sep 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wsmoses Sep 24, 2025

Choose a reason for hiding this comment

Uh oh!

wsmoses Sep 24, 2025

Choose a reason for hiding this comment

Uh oh!

wsmoses Sep 18, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Nov 29, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Nov 29, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Nov 29, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Nov 29, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Nov 29, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Nov 29, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Nov 29, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Nov 29, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Nov 29, 2025

Choose a reason for hiding this comment

Uh oh!

This comment was marked as outdated.

Uh oh!

jumerckx commented Nov 30, 2025

Uh oh!

wsmoses Dec 1, 2025

Choose a reason for hiding this comment

Uh oh!

wsmoses commented Dec 1, 2025

Uh oh!

jumerckx commented Dec 1, 2025

Uh oh!

giordano commented Dec 2, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

wsmoses Sep 24, 2025 •

edited

Loading