Separate Reactant number types by oschulz · Pull Request #2170 · EnzymeAD/Reactant.jl

oschulz · 2026-01-21T12:07:01Z

Addresses #1570.

Built on top of #3034.

oschulz · 2026-01-21T12:09:28Z

Let's see if we can make VS-Code-CoPilot/Claude do the heavy refactoring this will need and then fine-tune manually. ~~Claude Opus 4.5 was told ...~~ (obsolete, PR content now completely redone).

oschulz · 2026-01-21T12:16:02Z

~~Claude agent managed ...~~ (obsolete)

oschulz · 2026-01-21T12:28:25Z

~~Claude agent now came up with ...~ (obsolete)

oschulz · 2026-01-21T21:35:26Z

@wsmoses could you take a cursory initial look at whether this is going in the right direction, from your side?

It's not done yet - some AI-generated changes clearly will need manual polish, and I doubt all tests are passing at this point. But some stuff is working, and I think the bulk of the changes are probably more or less in by now.

I guess we want to move quickly with this, because there's high potential for conflicts with other incoming PRs, so this may diverge quickly from master.

Claude's summary of changes:

The refactoring successfully separated TracedRNumber and ConcreteRNumber types into specific Real (Integer/Float) and Complex hierarchies.

Key Changes Made:

Types.jl - New abstract and concrete type hierarchy:
- Added abstract types: RInteger{T} <: Integer, RFloat{T} <: AbstractFloat, RComplex{T} <: Number
- Added traced types: TracedRInteger{T}, TracedRFloat{T}, TracedRComplex{T}
- Added concrete PJRT types: ConcretePJRTInteger{T,D}, ConcretePJRTFloat{T,D}, ConcretePJRTComplex{T,D}
- Added concrete IFRT types: ConcreteIFRTInteger{T}, ConcreteIFRTFloat{T}, ConcreteIFRTComplex{T}
- Created union type aliases for backward compatibility: TracedRNumber{T}, ConcretePJRTNumber{T,D}, ConcreteIFRTNumber{T}
- Added helper functions _ConcretePJRTNumber and _ConcreteIFRTNumber for type dispatch
Reactant.jl - Updated exports and utilities:
- Added exports for new types: RInteger, RFloat, RComplex, TracedRInteger, TracedRFloat, TracedRComplex, etc.
- Extended unwrapped_eltype for new types
TracedRNumber.jl - Updated traced number operations:
- Added promotion rules for TracedRInteger, TracedRFloat, TracedRComplex
- Added convert methods for specific types
- Updated show methods for specific types
TracedPromotion.jl: Added promote_to for specific types
Tracing.jl: Updated tracing infrastructure:
- Extended is_traced_number to recognize specific types
- Updated traced_type_inner to return correct specific types
- Updated make_tracer to create correct specific types
Ops.jl: Added convert methods for specific types
Compiler.jl - Updated code generation:
- Fixed create_result to use correct specific types instead of union aliases.

wsmoses · 2026-01-22T04:52:21Z


-Base.zero(::TracedRNumber{T}) where {T} = Reactant.promote_to(TracedRNumber{T}, zero(T))
-Base.one(::TracedRNumber{T}) where {T} = Reactant.promote_to(TracedRNumber{T}, one(T))
+Base.zero(::TracedRInteger{T}) where {T} = Reactant.promote_to(TracedRInteger{T}, zero(T))


this is a lot of duplication...

Yes, there's a jungle of promote rules now that I definitely want to thin out.

How about overloading eltype to return the T out of TracedR* and ConcreteR*? That way you could just do:

Base.zero(::T) where {T<:RNumber} = Reactant.promote_to(T, zero(eltype(T)))

I would also include a method for zero(::Type{<:RNumber}) since that method also exists for regular Julia types.

actually, no need for eltype, just do

Base.zero(::T) where {G,T<:RNumber{G}} = Reactant.promote_to(T, zero(G))

Suggestions are implemented - I'll see if I can slim things down further.

oschulz · 2026-01-22T08:27:25Z

Thanks for the initial review @wsmoses ! I'll start do to polish then, and take a first look at the test failures.

mofeing

would you mind adding some tests to check that subtype hierarchy is preserved?

mofeing · 2026-01-23T07:35:43Z


-Base.zero(::TracedRNumber{T}) where {T} = Reactant.promote_to(TracedRNumber{T}, zero(T))
-Base.one(::TracedRNumber{T}) where {T} = Reactant.promote_to(TracedRNumber{T}, one(T))
+Base.zero(::TracedRInteger{T}) where {T} = Reactant.promote_to(TracedRInteger{T}, zero(T))


How about overloading eltype to return the T out of TracedR* and ConcreteR*? That way you could just do:

Base.zero(::T) where {T<:RNumber} = Reactant.promote_to(T, zero(eltype(T)))

I would also include a method for zero(::Type{<:RNumber}) since that method also exists for regular Julia types.

mofeing · 2026-01-23T07:40:16Z


-Base.zero(::TracedRNumber{T}) where {T} = Reactant.promote_to(TracedRNumber{T}, zero(T))
-Base.one(::TracedRNumber{T}) where {T} = Reactant.promote_to(TracedRNumber{T}, one(T))
+Base.zero(::TracedRInteger{T}) where {T} = Reactant.promote_to(TracedRInteger{T}, zero(T))


actually, no need for eltype, just do

Base.zero(::T) where {G,T<:RNumber{G}} = Reactant.promote_to(T, zero(G))

oschulz · 2026-01-24T11:46:24Z

would you mind adding some tests to check that subtype hierarchy is preserved?

Yes, I'll do that.

oschulz · 2026-01-25T22:37:42Z

Update: Had to add a type parameter RT to TracedArray, so we change from

TracedRArray{T,N} <: RArray{TracedRNumber{T},N}

to

TracedRArray{T,N,RT<:TracedRNumber{T}} <: RArray{RT,N}

It's seems the cleanest way to get the right element type into the RArray supertype.

Quite a few tests pass now, but a lot fail still.

oschulz · 2026-01-25T22:39:41Z

Question: Should CuTracedRNumber be split into separate number types as well?

oschulz · 2026-01-27T07:55:01Z

Currently figuring out some dispatch ambiguity problems.

oschulz · 2026-07-08T14:05:36Z

@wsmoses @avik-pal I've replaced this with a completely fresh attempt, using a lot of Claude Fable at max effort. And I think this time it succeeded.

This needs EnzymeAD/Enzyme.jl#3293.

oschulz · 2026-07-08T21:56:35Z

This PR is on top of #3034

Loosen invariant `TracedRNumber{T}`/`TracedRArray{T,N}` type parameters to covariant `<:` bounds in method signatures and type aliases. These are no-ops today (the types are concrete) but make the signatures robust to element types that subtype the traced numbers. The `AnyTracedRArray` alias itself stays invariant: a covariant bound would match `Vector{Union{}}` (the type of empty array literals), which must not be routed into the traced-array machinery (EnzymeAD#1290). Created by generative AI.

…ed array `similar(::ConcreteArray, ::Type{S}, dims)` passed `S` straight to the device buffer, whose `primitive_type(S)` is undefined for wrapped Reactant number types. This surfaced as a `MethodError` from LinearAlgebra's generic `\` (and other `similar`-based allocations) when a concrete array meets a traced result element type during tracing, e.g. `UpperTriangular(m) \ b` with a captured concrete `b`. A concrete buffer cannot hold traced values, so route such requests to a `TracedRArray` of the unwrapped element type. Created by generative AI.

Traced and concrete Reactant numbers are now split into integer, float, and complex types so that they subtype Integer and AbstractFloat where possible (issue EnzymeAD#1570). Julia has no abstract complex type, so complex Reactant numbers remain subtypes of Number. RNumber, TracedRNumber, AbstractConcreteNumber, ConcretePJRTNumber, and ConcreteIFRTNumber become union aliases over the new kind-specific types, with pseudo-constructors keeping the existing constructor API working. TracedRArray gains a third type parameter carrying the traced element type, so that e.g. a traced float array is an AbstractArray{<:AbstractFloat}. Created by generative AI.

Promotion collapses to a single parametric rule since promoting the primitive element types already selects the numeric kind. Operations move from the TracedRNumber union to the most specific honest kind sets (TracedRReal, TracedRInteger, TracedRFloat, TracedRComplex) so they dominate Base's methods on Real, Integer, and AbstractFloat, with explicit disambiguation against Base's specialized comparison methods. Base identities (real/imag/conj on Real, isinteger, isfinite on Integer) now apply directly and replace the corresponding overloads. Created by generative AI.

traced_type_inner returns the kind-specific concrete types and handles the union aliases and the TracedRArray UnionAll in type introspection. Kind-specific methods resolve the ambiguities between the identity methods on Integer/AbstractFloat/RNumber and the traced/concrete number methods. Also fixes a latent typo (ConcretePJRTNumbe) in the ArrayToConcrete branch. Created by generative AI.

…ypes Type-level signatures widen to subtype form, invariant traced eltypes in wrapper types become covariant, and mlir_type learns the two-parameter TracedRArray UnionAll. Static-vs-traced integer checks in top_k and standardize_start_index now exclude traced integers explicitly. Colon overloads move to TracedRReal so they dominate Base's promoting colon. TracedRational simplifies: RInteger <: Integer covers traced and concrete integers, collapsing the // constructor product loop. Created by generative AI.

Concrete number operations are defined per numeric kind with disambiguation against Base's specialized methods, mirroring the traced side. Base.decompose forwards to the host value so AbstractFloat hashing and printing work. create_result gains per-kind methods to avoid ambiguity with the Union{Integer,AbstractFloat,...} method. Enzyme activity is pinned to Duplicated for Reactant numbers, which would otherwise be classified Active now that they subtype AbstractFloat. Created by generative AI.

Created by generative AI.

Disambiguation follows Base's exact method groupings: single rounding modes plus the grouped nearest-modes union for div, concrete float types for copysign and ^, per-type float8 promotion rules, and concrete-integer variants of the TracedRational operations. Traced integer range indices dispatch via TracedRInteger. Created by generative AI.

Adds value constructors for the kind-specific concrete number types, order comparisons for concrete numbers (now Real, so Base's same-type fallback throws), real-only diagonals for rem/mod/mod1 that dominate the mixed-argument methods, isreal for traced complex numbers (Base has no Number fallback), and sincos/exp2/exp10/mod2pi/modf for traced floats where Base's f(x::Real) = f(float(x)) pattern would recurse forever. Empty (Union{}-eltype) arrays are guarded from the covariant AnyTracedRArray alias, and apply_type_with_promotion concretizes the TracedRArray UnionAll before applying type parameters. Created by generative AI.

…tization Base's ^(::Number, ::Integer) calls power_by_squaring, which branches on traced booleans when the exponent is a traced integer; ^ with traced integer exponents now promotes and lowers to stablehlo.power. Enzyme's gradient sugar classifies traced floats as Active scalars (its AbstractFloat fast path runs in a fixed world, so overriding active_reg_nothrow cannot help); overload_autodiff now supports Active traced scalars by seeding re-emitted arguments with a zero cotangent and returning gradients as fresh traced values. traced_type now produces the concrete three-parameter TracedRArray so wrapper structs with eltype-constrained parameters instantiate. Broadcast eltype recovery triggers for any unusable inferred eltype (e.g. Real), and fma on traced floats avoids Base's same-type error path. Created by generative AI.

Created by generative AI.

CuTracedRInteger/CuTracedRFloat/CuTracedRComplex mirror the traced host types so that kernel-side numbers subtype Integer and AbstractFloat, with CuTracedRNumber as the union alias and pseudo-constructors keeping the existing API. Kernel-side operations follow the same per-kind dispatch scheme as the host types. Created by generative AI.

Traced reals now take AbstractFFTs' Real conversion path, so the FFT extension lowers fftfloat/complexfloat/realfloat directly instead of copying through host arrays. TensorOperations allocation returns the concrete traced array type, which also gains undef constructors in the three-parameter form. The TracedRArray inner constructor now rejects an element type inconsistent with the primitive type, and the remaining constructor/^ ambiguities against Base's BigFloat and Rational methods are resolved. Created by generative AI.

Enzyme now recognizes the kind-split Reactant number types ahead of its raw-float fast path (Enzyme.jl branch fixes-for-reactant), so the `active_reg_nothrow` and `guess_activity` overrides are obsolete; the native classification is also more precise (`Const` for integer kinds, mode-dependent annotations). Explicit `Active` scalar arguments keep working through `overload_autodiff` and are now tested directly, since the `gradient` sugar takes the `Duplicated` route again. Created by generative AI.

The full ambiguity scan of the qa suite surfaced pairs beyond the earlier targeted checks. Two rules recur: disambiguators must mirror the exact TypeVar shape of the competing Base method, and `promote_type` evaluates `promote_rule` in both argument orders, so special-cased rules need reverse-direction twins. - per-kind `BigInt`/`BigFloat` promotion rules and same-`T` diagonals for `^` and the nearest-mode `div` methods - `+`, `-`, `*` between traced integers and `Rational` promote to `TracedRational` instead of hitting Base's checked integer arithmetic; `Rational^tracedint` computes in floating point - consistent traced-range `getindex` grid; drop a redundant duplicate of the LinRange/StepRangeLen method - `searchsortedfirst` on ranges promotes the needle jointly with the range element type instead of truncating it - sparse and `ReadOnly` `getindex` disambiguators in the SparseArrays extension - reference types directly in the `create_result` loop and drop a stale mid-file import, per ExplicitImports Created by generative AI.

Created by generative AI.

A formatter-version drift during development left changes on code the split does not touch. Restore them to match main: - ext/ReactantKernelAbstractionsExt.jl: the `VERSION < 1.12` kernel-call overload had been mangled from a `(obj)(args...)` functor into a plain one-arg anonymous function (the arg list became a body statement), silently dropping the overload on Julia < 1.12. Restored. - ext/ReactantMPIExt/Ops.jl: revert stray `return throw(...)`. - redundant interpolation parentheses (`=($(x))` -> `=$(x)`) in the FFT, LinearAlgebra, and CUDA sources. Created by generative AI.

oschulz · 2026-07-09T07:54:47Z

I'm working on the test failures here and in #3034.

With the number kinds in place, `TracedRNumber{T}` is a finite union of concrete types, so `AnyTracedRArray` can enumerate the traced element types invariantly instead of bounding them covariantly. A covariant bound (`AbstractArray{<:TracedRNumber{T},N}`) would also match `Vector{Union{}}` (the type of empty array literals), which must not be routed into the traced-array machinery (EnzymeAD#1290). This restores the upstream subtype guard test alongside the behavior tests. Created by generative AI.

`traced_type` now yields `TracedRArray{T,N,RT}` with the concrete traced element type filled in; spell that out in the `trace_type` expectation table. Created by generative AI.

`concretize_traced_param` had been inserted between the docstring and `apply_type_with_promotion`, detaching the docstring and breaking the documentation build. Also correct the docstring example: the params are passed as a vector, and the traced array parameter now comes back kind-typed (and flagged as changed). Created by generative AI.

…ping Julia 1.10's `getindex(::UnitRange{<:Union{Bool,Int8,...}}, ::Integer)` became ambiguous with the traced-integer range indexing once traced integers subtype `Integer`; add the exact-grouping disambiguator. Created by generative AI.

Enzyme classifies wrapped-number activity by type name; 0.13.177 is the first release that recognizes the kind-split traced and concrete number types before the `AbstractFloat` activity fast path. Created by generative AI.

oschulz marked this pull request as draft January 21, 2026 12:07