From 5bbfb25e70177fe0b2be05e8e5ed4a34f756a7b2 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Wed, 8 Jul 2026 20:47:50 +0200
Subject: [PATCH 01/30] Widen traced-type signatures to accept covariant
 element types

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.

Deliberately left invariant:
- the `AnyTracedRArray` alias: a covariant bound would match
  `Vector{Union{}}` (the type of empty array literals), which must not be
  routed into the traced-array machinery (#1290);
- the `promote_rule` methods against `Missing`/`Nothing` and the
  `LinRange`/`StepRangeLen` `getindex` resolver: these are exact-shape
  disambiguators, and loosening them makes their specificity against
  their Base counterparts undecidable (method ambiguities).

Created by generative AI.
---
 ext/ReactantAbstractFFTsExt.jl          |  2 +-
 ext/ReactantCUDAExt.jl                  |  4 +--
 ext/ReactantSparseArraysExt/Errors.jl   |  6 ++--
 ext/ReactantSparseArraysExt/ReadOnly.jl | 34 ++++++++++-----------
 lib/ReactantCore/src/ReactantCore.jl    |  2 +-
 src/ConcreteRArray.jl                   |  8 ++---
 src/Indexing.jl                         | 14 ++++-----
 src/Ops.jl                              | 22 +++++++-------
 src/TracedPromotion.jl                  | 32 ++++++++++----------
 src/TracedRArray.jl                     |  4 +--
 src/TracedRNumber.jl                    | 12 ++++----
 src/TracedRange.jl                      |  2 +-
 src/TracedUtils.jl                      |  8 ++---
 src/stdlibs/LinearAlgebra.jl            | 40 ++++++++++++-------------
 14 files changed, 95 insertions(+), 95 deletions(-)

diff --git a/ext/ReactantAbstractFFTsExt.jl b/ext/ReactantAbstractFFTsExt.jl
index c29bbad9fb..3b5aeff7ee 100644
--- a/ext/ReactantAbstractFFTsExt.jl
+++ b/ext/ReactantAbstractFFTsExt.jl
@@ -7,7 +7,7 @@ using Reactant.Ops: @opcall
 
 # FFTW.jl defines methods on StridedArrays, and we have to be more specific than FFTW to
 # catch its methods.
-const AnyStridedTracedRArray{T,N} = StridedArray{TracedRNumber{T},N}
+const AnyStridedTracedRArray{T,N} = StridedArray{<:TracedRNumber{T},N}
 
 # To automatically convert FFT plans to traced versions
 # To extend a user needs to extend Reactant.reactant_fftplan for their plan type
diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl
index 43d322984f..35ea0098a2 100644
--- a/ext/ReactantCUDAExt.jl
+++ b/ext/ReactantCUDAExt.jl
@@ -144,7 +144,7 @@ Base.OneTo(x::CuTracedRNumber{<:Integer}) = Base.OneTo(x[])
     end
 end
 
-@inline function Base.convert(CT::Type{CuTracedRNumber{Float64,1}}, x::Number)
+@inline function Base.convert(CT::Type{<:CuTracedRNumber{Float64,1}}, x::Number)
     return CT(
         Base.reinterpret(
             Core.LLVMPtr{Float64,1},
@@ -168,7 +168,7 @@ end
     )
 end
 
-@inline function Base.convert(CT::Type{CuTracedRNumber{Float32,1}}, x::Number)
+@inline function Base.convert(CT::Type{<:CuTracedRNumber{Float32,1}}, x::Number)
     return CT(
         Base.reinterpret(
             Core.LLVMPtr{Float32,1},
diff --git a/ext/ReactantSparseArraysExt/Errors.jl b/ext/ReactantSparseArraysExt/Errors.jl
index 23622e62cd..5f8de6b426 100644
--- a/ext/ReactantSparseArraysExt/Errors.jl
+++ b/ext/ReactantSparseArraysExt/Errors.jl
@@ -1,15 +1,15 @@
 # We don't yet support sparse arrays. Throwing errors on this dispatches to avoid
 # ambiguities.
-function Base.getindex(::AbstractSparseArray{TracedRNumber{T},N,1}, ::Int64) where {T,N}
+function Base.getindex(::AbstractSparseArray{<:TracedRNumber{T},N,1}, ::Int64) where {T,N}
     return error("Sparse arrays are not supported by reactant yet.")
 end
 
 function Base.getindex(
-    ::AbstractSparseMatrixCSC{TracedRNumber{T}}, ::Int64, ::Int64
+    ::AbstractSparseMatrixCSC{<:TracedRNumber{T}}, ::Int64, ::Int64
 ) where {T}
     return error("Sparse arrays are not supported by reactant yet.")
 end
 
-function Base.getindex(::CHOLMOD.Sparse{TracedRNumber{T}}, ::Int64, ::Int64) where {T}
+function Base.getindex(::CHOLMOD.Sparse{<:TracedRNumber{T}}, ::Int64, ::Int64) where {T}
     return error("Sparse arrays are not supported by reactant yet.")
 end
diff --git a/ext/ReactantSparseArraysExt/ReadOnly.jl b/ext/ReactantSparseArraysExt/ReadOnly.jl
index 90b3bb7f9d..6c8630300c 100644
--- a/ext/ReactantSparseArraysExt/ReadOnly.jl
+++ b/ext/ReactantSparseArraysExt/ReadOnly.jl
@@ -1,48 +1,48 @@
 function Base.getindex(
-    A::ReadOnly{Reactant.TracedRNumber{T},1,V}, inds::AbstractArray
-) where {T,V<:AbstractArray{Reactant.TracedRNumber{T},1}}
+    A::ReadOnly{<:Reactant.TracedRNumber{T},1,V}, inds::AbstractArray
+) where {T,V<:AbstractArray{<:Reactant.TracedRNumber{T},1}}
     return getindex(parent(A), inds)
 end
 
 function Base.getindex(
-    A::ReadOnly{Reactant.TracedRNumber{T},1,V}, ::Colon
-) where {T,V<:AbstractArray{Reactant.TracedRNumber{T},1}}
+    A::ReadOnly{<:Reactant.TracedRNumber{T},1,V}, ::Colon
+) where {T,V<:AbstractArray{<:Reactant.TracedRNumber{T},1}}
     return getindex(parent(A), :)
 end
 
 function Base.getindex(
-    A::ReadOnly{Reactant.TracedRNumber{T},N,V},
+    A::ReadOnly{<:Reactant.TracedRNumber{T},N,V},
     i::Union{Integer,Reactant.TracedRNumber{<:Integer}},
-) where {T,N,V<:AbstractArray{Reactant.TracedRNumber{T},N}}
+) where {T,N,V<:AbstractArray{<:Reactant.TracedRNumber{T},N}}
     return getindex(parent(A), i)
 end
 
 function Base.getindex(
-    A::ReadOnly{Reactant.TracedRNumber{T},N,V}, inds::AbstractArray
-) where {T,N,V<:AbstractArray{Reactant.TracedRNumber{T},N}}
+    A::ReadOnly{<:Reactant.TracedRNumber{T},N,V}, inds::AbstractArray
+) where {T,N,V<:AbstractArray{<:Reactant.TracedRNumber{T},N}}
     return getindex(parent(A), inds)
 end
 
 function Base.getindex(
-    A::ReadOnly{Reactant.TracedRNumber{T},N,V}, idx::CartesianIndex{N}
-) where {T,N,V<:AbstractArray{Reactant.TracedRNumber{T},N}}
+    A::ReadOnly{<:Reactant.TracedRNumber{T},N,V}, idx::CartesianIndex{N}
+) where {T,N,V<:AbstractArray{<:Reactant.TracedRNumber{T},N}}
     return getindex(parent(A), idx)
 end
 
 function Base.getindex(
-    A::ReadOnly{Reactant.TracedRNumber{T},1,V}, idx::CartesianIndex{1}
-) where {T,V<:AbstractArray{Reactant.TracedRNumber{T},1}}
+    A::ReadOnly{<:Reactant.TracedRNumber{T},1,V}, idx::CartesianIndex{1}
+) where {T,V<:AbstractArray{<:Reactant.TracedRNumber{T},1}}
     return getindex(parent(A), idx)
 end
 
 function Base.getindex(
-    A::ReadOnly{Reactant.TracedRNumber{T},N,V}, ::Colon
-) where {T,N,V<:AbstractArray{Reactant.TracedRNumber{T},N}}
+    A::ReadOnly{<:Reactant.TracedRNumber{T},N,V}, ::Colon
+) where {T,N,V<:AbstractArray{<:Reactant.TracedRNumber{T},N}}
     return getindex(parent(A), :)
 end
 
 function Base.getindex(
-    x::ReadOnly{TracedRNumber{T},N},
+    x::ReadOnly{<:TracedRNumber{T},N},
     idx::Vararg{Union{Integer,Reactant.TracedRNumber{<:Integer}},N},
 ) where {T,N}
     return getindex(parent(x), idx...)
@@ -51,14 +51,14 @@ end
 function Base.getindex(
     x::SparseArrays.ReadOnly{Reactant.TracedRNumber{T},1,V},
     idx::Union{Int64,Reactant.TracedRNumber{Int64}},
-) where {T,V<:AbstractArray{Reactant.TracedRNumber{T},1}}
+) where {T,V<:AbstractArray{<:Reactant.TracedRNumber{T},1}}
     return getindex(parent(x), idx)
 end
 
 function Base.getindex(
     x::SparseArrays.ReadOnly{
         Reactant.TracedRNumber{T},N,V
-    } where {V<:AbstractArray{Reactant.TracedRNumber{T},N}},
+    } where {V<:AbstractArray{<:Reactant.TracedRNumber{T},N}},
     idx::Vararg{Union{Int64,Reactant.TracedRNumber{Int64}},N},
 ) where {T,N}
     return getindex(parent(x), idx...)
diff --git a/lib/ReactantCore/src/ReactantCore.jl b/lib/ReactantCore/src/ReactantCore.jl
index 17aeb3fc92..18e2ea26ea 100644
--- a/lib/ReactantCore/src/ReactantCore.jl
+++ b/lib/ReactantCore/src/ReactantCore.jl
@@ -886,7 +886,7 @@ end
 """
     materialize_traced_array(AbstractArray{<:TracedRNumber})::TracedRArray
 
-Given an AbstractArray{TracedRNumber}, return or create an equivalent TracedRArray.
+Given an AbstractArray{<:TracedRNumber}, return or create an equivalent TracedRArray.
 
 """
 function materialize_traced_array end
diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 414ea4222c..c807ba9ab0 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -725,8 +725,8 @@ for aType in (:ConcretePJRTArray, :ConcreteIFRTArray)
         end
 
         function Base.copyto!(
-            dest::SubArray{TracedRNumber{T1},<:Any,<:$(aType)},
-            src::SubArray{TracedRNumber{T2},<:Any,<:Array},
+            dest::SubArray{<:TracedRNumber{T1},<:Any,<:$(aType)},
+            src::SubArray{<:TracedRNumber{T2},<:Any,<:Array},
         ) where {T1,T2}
             throw(MethodError(copyto!, (dest, src)))
         end
@@ -771,7 +771,7 @@ function Base.zero(x::ConcreteIFRTArray{T,N}) where {T,N}
 end
 
 function Base.fill!(
-    a::ConcretePJRTArray{TracedRNumber{T},N}, val::TracedRNumber{T2}
+    a::ConcretePJRTArray{<:TracedRNumber{T},N}, val::TracedRNumber{T2}
 ) where {T,T2,N}
     throw(MethodError(fill!, (a, val)))
 end
@@ -805,7 +805,7 @@ function Base.fill!(a::ConcreteIFRTArray{T,N}, val) where {T,N}
 end
 
 function Base.fill!(
-    a::ConcreteIFRTArray{TracedRNumber{T},N}, val::TracedRNumber{T2}
+    a::ConcreteIFRTArray{<:TracedRNumber{T},N}, val::TracedRNumber{T2}
 ) where {T,T2,N}
     throw(MethodError(fill!, (a, val)))
 end
diff --git a/src/Indexing.jl b/src/Indexing.jl
index 7768a1448b..e3f664baf9 100644
--- a/src/Indexing.jl
+++ b/src/Indexing.jl
@@ -77,7 +77,7 @@ end
 
 # This method is needed exclusively to resolve an ambiguity
 function Base.getindex(
-    l::Selectors.List{TracedRNumber{T}}, index::Union{Int,TracedRNumber{Int}}
+    l::Selectors.List{<:TracedRNumber{T}}, index::Union{Int,TracedRNumber{Int}}
 ) where {T}
     return getindex(Reactant.promote_to(TracedRArray{T,1}, l), index)
 end
@@ -173,8 +173,8 @@ end
 
 ### Specialize certain dispatches for better codegen
 for aType in (
-    Base.ReshapedArray{TracedRNumber{T}} where {T},
-    PermutedDimsArray{TracedRNumber{T}} where {T},
+    Base.ReshapedArray{<:TracedRNumber{T}} where {T},
+    PermutedDimsArray{<:TracedRNumber{T}} where {T},
 )
     @eval begin
         function Base.getindex(a::$(aType), indices::Union{Int,TracedRNumber{Int}}...)
@@ -188,8 +188,8 @@ for aType in (
 end
 
 for aType in (
-    Base.ReshapedArray{TracedRNumber{T},N,P,Tuple{}} where {T,N,P<:AbstractArray},
-    Base.ReshapedArray{TracedRNumber{T},1,P,Tuple{}} where {T,P<:AbstractArray},
+    Base.ReshapedArray{<:TracedRNumber{T},N,P,Tuple{}} where {T,N,P<:AbstractArray},
+    Base.ReshapedArray{<:TracedRNumber{T},1,P,Tuple{}} where {T,P<:AbstractArray},
 )
     @eval function Base.getindex(a::$(aType), indices::Int)
         return getindex(materialize_traced_array(a), indices)
@@ -197,7 +197,7 @@ for aType in (
 end
 
 function Base.getindex(
-    x::Base.ReshapedArray{TracedRNumber{T}}, index::Base.ReshapedIndex
+    x::Base.ReshapedArray{<:TracedRNumber{T}}, index::Base.ReshapedIndex
 ) where {T}
     return getindex(parent(x), index.parentindex)
 end
@@ -463,7 +463,7 @@ function generate_index_list(i1, is...)
 end
 
 function scalar_index_to_cartesian(
-    idx::AbstractVector{TracedRNumber{T}}, sz::NTuple{N,Int}
+    idx::AbstractVector{<:TracedRNumber{T}}, sz::NTuple{N,Int}
 ) where {T,N}
     idx = materialize_traced_array(idx)
     idx = @opcall(subtract(idx, @opcall(fill(T(1), size(idx)))))
diff --git a/src/Ops.jl b/src/Ops.jl
index b63482680e..883cce7b46 100644
--- a/src/Ops.jl
+++ b/src/Ops.jl
@@ -355,7 +355,7 @@ function _fill_element_attr(x::Complex)
 end
 
 @noinline function concatenate(
-    inputs::Vector{TracedRArray{T,N}},
+    inputs::Vector{<:TracedRArray{T,N}},
     dimension::Int;
     location=mlir_stacktrace("fill", @__FILE__, @__LINE__),
 ) where {T,N}
@@ -869,7 +869,7 @@ end
 end
 
 function bitcast_convert(
-    ::Type{TracedRArray{U,N}},
+    ::Type{<:TracedRArray{U,N}},
     x::TracedRArray{T,N};
     location=mlir_stacktrace("bitcast_convert", @__FILE__, @__LINE__),
 ) where {T,U,N}
@@ -1738,7 +1738,7 @@ end
 end
 
 @noinline function rng_bit_generator(
-    ::Type{TracedRNumber{T}}, seed::TracedRArray{UInt64,1}, shape; kwargs...
+    ::Type{<:TracedRNumber{T}}, seed::TracedRArray{UInt64,1}, shape; kwargs...
 ) where {T}
     return rng_bit_generator(T, seed, shape; kwargs...)
 end
@@ -1798,7 +1798,7 @@ end
 end
 
 @noinline function randn(
-    ::Type{TracedRNumber{T}}, seed::TracedRArray{UInt64,1}, shape; kwargs...
+    ::Type{<:TracedRNumber{T}}, seed::TracedRArray{UInt64,1}, shape; kwargs...
 ) where {T}
     return randn(T, seed, shape; kwargs...)
 end
@@ -1841,7 +1841,7 @@ distribution with rate 1. Returns a NamedTuple with the following fields:
 end
 
 @noinline function randexp(
-    ::Type{TracedRNumber{T}}, seed::TracedRArray{UInt64,1}, shape; kwargs...
+    ::Type{<:TracedRNumber{T}}, seed::TracedRArray{UInt64,1}, shape; kwargs...
 ) where {T}
     return randexp(T, seed, shape; kwargs...)
 end
@@ -1923,7 +1923,7 @@ end
 
 # eltype conversion
 @noinline function convert(
-    ::Type{TracedRArray{T,N}},
+    ::Type{<:TracedRArray{T,N}},
     x::TracedRArray;
     location=mlir_stacktrace("convert", @__FILE__, @__LINE__),
 ) where {T,N}
@@ -1940,7 +1940,7 @@ end
 end
 
 @noinline function convert(
-    ::Type{TracedRNumber{T}},
+    ::Type{<:TracedRNumber{T}},
     x::TracedRNumber;
     location=mlir_stacktrace("convert", @__FILE__, @__LINE__),
 ) where {T}
@@ -2166,7 +2166,7 @@ end
 
 @noinline function scatter(
     f::F,
-    dest::Vector{TracedRArray{T,N}},
+    dest::Vector{<:TracedRArray{T,N}},
     scatter_indices::TracedRArray{Int64},
     updates::Vector{<:TracedRArray{T}};
     location=mlir_stacktrace("scatter", @__FILE__, @__LINE__),
@@ -2195,7 +2195,7 @@ end
 end
 
 @noinline function scatter(
-    dest::Vector{TracedRArray{T,N}},
+    dest::Vector{<:TracedRArray{T,N}},
     scatter_indices::TracedRArray{TI},
     updates::Vector{<:TracedRArray{T}};
     update_computation::MLIR.IR.Region,
@@ -3967,8 +3967,8 @@ end
 
 @noinline function reduce_window(
     f::F,
-    inputs::Vector{TracedRArray{T,N}},
-    init_values::Vector{TracedRNumber{T}};
+    inputs::Vector{<:TracedRArray{T,N}},
+    init_values::Vector{<:TracedRNumber{T}};
     window_dimensions::Vector{Int},
     window_strides::Vector{Int},
     base_dilations::Vector{Int},
diff --git a/src/TracedPromotion.jl b/src/TracedPromotion.jl
index 765f26fa83..00acc7934f 100644
--- a/src/TracedPromotion.jl
+++ b/src/TracedPromotion.jl
@@ -11,30 +11,30 @@ function promote_to(::Type{TracedRArray{T}}, rhs) where {T}
     return promote_to(TracedRArray{T,ndims(rhs)}, rhs)
 end
 
-promote_to(::Type{TracedRArray{T,N}}, rhs::TracedRArray{T,N}) where {T,N} = rhs
-function promote_to(::Type{TracedRArray{T,N}}, rhs::TracedRArray{T2,N}) where {T,T2,N}
+promote_to(::Type{<:TracedRArray{T,N}}, rhs::TracedRArray{T,N}) where {T,N} = rhs
+function promote_to(::Type{<:TracedRArray{T,N}}, rhs::TracedRArray{T2,N}) where {T,T2,N}
     return @opcall convert(TracedRArray{T,N}, rhs)
 end
 
 function promote_to(
-    ::Type{TracedRArray{T,N}}, rhs::AbstractArray{<:TracedRNumber,N}
+    ::Type{<:TracedRArray{T,N}}, rhs::AbstractArray{<:TracedRNumber,N}
 ) where {T,N}
     return @opcall convert(TracedRArray{T,N}, aos_to_soa(materialize_traced_array(rhs)))
 end
 
-function promote_to(::Type{TracedRArray{T,1}}, rhs::AbstractRange) where {T}
+function promote_to(::Type{<:TracedRArray{T,1}}, rhs::AbstractRange) where {T}
     return @opcall convert(TracedRArray{T,1}, materialize_traced_array(rhs))
 end
 
-function promote_to(::Type{TracedRArray{T,1}}, rhs::Base.OneTo) where {T}
+function promote_to(::Type{<:TracedRArray{T,1}}, rhs::Base.OneTo) where {T}
     return promote_to(TracedRArray{T,1}, first(rhs):last(rhs))
 end
 
-function promote_to(::Type{TracedRArray{T,N}}, rhs::LinearAlgebra.Diagonal) where {T,N}
+function promote_to(::Type{<:TracedRArray{T,N}}, rhs::LinearAlgebra.Diagonal) where {T,N}
     return LinearAlgebra.diagm(promote_to(TracedRArray{T,1}, rhs.diag))
 end
 
-function promote_to(::Type{TracedRArray{T,N}}, rhs::AbstractArray{<:Any,N}) where {T,N}
+function promote_to(::Type{<:TracedRArray{T,N}}, rhs::AbstractArray{<:Any,N}) where {T,N}
     if ancestor(rhs) isa AnyTracedRArray
         return promote_to(TracedRArray{T,N}, materialize_traced_array(rhs))
     end
@@ -42,7 +42,7 @@ function promote_to(::Type{TracedRArray{T,N}}, rhs::AbstractArray{<:Any,N}) wher
     return promote_to(TracedRArray{T,N}, @opcall(constant(rhs)))
 end
 
-function promote_to(::Type{TracedRArray{T,N}}, rhs::Number) where {T,N}
+function promote_to(::Type{<:TracedRArray{T,N}}, rhs::Number) where {T,N}
     return promote_to(TracedRArray{T,N}, [rhs])
 end
 
@@ -58,35 +58,35 @@ function promote_to(
     return promote_to(TracedRNumber{float(unwrapped_eltype(T))}, rhs)
 end
 
-promote_to(::Type{TracedRNumber{T}}, rhs::TracedRNumber{T}) where {T} = rhs
-function promote_to(::Type{TracedRNumber{T}}, rhs::TracedRNumber{T2}) where {T,T2}
+promote_to(::Type{<:TracedRNumber{T}}, rhs::TracedRNumber{T}) where {T} = rhs
+function promote_to(::Type{<:TracedRNumber{T}}, rhs::TracedRNumber{T2}) where {T,T2}
     return @opcall convert(TracedRNumber{T}, rhs)
 end
 function promote_to(
-    ::Type{TracedRNumber{T}}, rhs::TracedRNumber{T2}
+    ::Type{<:TracedRNumber{T}}, rhs::TracedRNumber{T2}
 ) where {T<:ReactantFloat8,T2}
     return @opcall convert(TracedRNumber{T}, rhs)
 end
 
-function promote_to(::Type{TracedRArray{T,0}}, rhs::TracedRNumber{T2}) where {T,T2}
+function promote_to(::Type{<:TracedRArray{T,0}}, rhs::TracedRNumber{T2}) where {T,T2}
     return TracedRArray{T,0}((), @opcall(convert(TracedRNumber{T}, rhs)).mlir_data, ())
 end
-function promote_to(::Type{TracedRNumber{T}}, rhs::TracedRArray{T2,0}) where {T,T2}
+function promote_to(::Type{<:TracedRNumber{T}}, rhs::TracedRArray{T2,0}) where {T,T2}
     return TracedRNumber{T}((), @opcall(convert(TracedRArray{T,0}, rhs)).mlir_data)
 end
 
-function promote_to(::Type{TracedRNumber{T}}, rhs::Number) where {T<:ReactantFloat8}
+function promote_to(::Type{<:TracedRNumber{T}}, rhs::Number) where {T<:ReactantFloat8}
     res = @opcall(fill(rhs))
     return @opcall convert(
         TracedRNumber{T}, TracedRNumber{unwrapped_eltype(res)}((), res.mlir_data)
     )
 end
 
-function promote_to(::Type{TracedRNumber{T}}, rhs::Number) where {T}
+function promote_to(::Type{<:TracedRNumber{T}}, rhs::Number) where {T}
     return promote_to(TracedRNumber{T}, T(rhs))
 end
 
-function promote_to(::Type{TracedRNumber{T}}, rhs::T) where {T<:ReactantPrimitive}
+function promote_to(::Type{<:TracedRNumber{T}}, rhs::T) where {T<:ReactantPrimitive}
     res = @opcall fill(rhs)
     return TracedRNumber{T}((), res.mlir_data)
 end
diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
index f5625fbc52..0d4144753c 100644
--- a/src/TracedRArray.jl
+++ b/src/TracedRArray.jl
@@ -23,7 +23,7 @@ Base.strides(x::TracedRArray) = Base.size_to_strides(1, size(x)...)
 
 Base.IndexStyle(::Type{<:TracedRArray}) = Base.IndexLinear()
 
-Base.elsize(::Type{TracedRArray{T,N}}) where {T,N} = sizeof(T)
+Base.elsize(::Type{<:TracedRArray{T,N}}) where {T,N} = sizeof(T)
 
 # This is required otherwise we will copy a tracedrarray each time
 # we use it
@@ -289,7 +289,7 @@ function Base.similar(
 end
 
 function Base.similar(
-    ::Broadcasted{AbstractReactantArrayStyle{N}}, ::Type{TracedRNumber{T}}, dims
+    ::Broadcasted{AbstractReactantArrayStyle{N}}, ::Type{<:TracedRNumber{T}}, dims
 ) where {T<:Reactant.ReactantPrimitive,N}
     @assert N isa Int
     return (@opcall fill(zero(T), dims))::TracedRArray{T,N}
diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index c3a1377390..8289603137 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -30,12 +30,12 @@ function Base.eps(::Type{TracedRNumber{T}}) where {T}
 end
 Base.eps(x::TracedRNumber{T}) where {T} = eps(typeof(x))
 
-function Base.typemin(::Type{TracedRNumber{T}}) where {T}
+function Base.typemin(::Type{<:TracedRNumber{T}}) where {T}
     return Reactant.promote_to(TracedRNumber{T}, typemin(T))
 end
 Base.typemin(x::TracedRNumber{T}) where {T} = typemin(typeof(x))
 
-function Base.typemax(::Type{TracedRNumber{T}}) where {T}
+function Base.typemax(::Type{<:TracedRNumber{T}}) where {T}
     return Reactant.promote_to(TracedRNumber{T}, typemax(T))
 end
 Base.typemax(x::TracedRNumber{T}) where {T} = typemax(typeof(x))
@@ -1053,17 +1053,17 @@ function Base.copysign(x::S, y::TracedRNumber{T}) where {S<:Number,T}
     return copysign(Reactant.promote_to(TracedRNumber{S}, x), y)
 end
 
-function Base.zeros(::Type{TracedRNumber{T}}, dims::Dims{N}) where {T,N}
+function Base.zeros(::Type{<:TracedRNumber{T}}, dims::Dims{N}) where {T,N}
     return @opcall fill(zero(T), dims)
 end
-function Base.zeros(::Type{TracedRNumber{T}}, ::Tuple{}) where {T}
+function Base.zeros(::Type{<:TracedRNumber{T}}, ::Tuple{}) where {T}
     return @opcall fill(zero(T), ())
 end
 
-function Base.ones(::Type{TracedRNumber{T}}, dims::Dims{N}) where {T,N}
+function Base.ones(::Type{<:TracedRNumber{T}}, dims::Dims{N}) where {T,N}
     return @opcall fill(one(T), dims)
 end
-function Base.ones(::Type{TracedRNumber{T}}, ::Tuple{}) where {T}
+function Base.ones(::Type{<:TracedRNumber{T}}, ::Tuple{}) where {T}
     return @opcall fill(one(T), ())
 end
 
diff --git a/src/TracedRange.jl b/src/TracedRange.jl
index 25ec0a98ac..1e8fa1ed93 100644
--- a/src/TracedRange.jl
+++ b/src/TracedRange.jl
@@ -154,7 +154,7 @@ function Base._in_unit_range(
     return (i > 0) & (val <= v.stop) & (val >= v.start)
 end
 
-@inline function Base.length(r::TracedUnitRange{TracedRNumber{T}}) where {T}
+@inline function Base.length(r::TracedUnitRange{<:TracedRNumber{T}}) where {T}
     if r.length >= 0
         return r.length
     end
diff --git a/src/TracedUtils.jl b/src/TracedUtils.jl
index ee74b8a727..a1bb430e32 100644
--- a/src/TracedUtils.jl
+++ b/src/TracedUtils.jl
@@ -107,7 +107,7 @@ function set_mlir_data!(x::TracedRArray, data)
     return x
 end
 
-function set_mlir_data!(x::Base.ReshapedArray{TracedRNumber{T}}, data) where {T}
+function set_mlir_data!(x::Base.ReshapedArray{<:TracedRNumber{T}}, data) where {T}
     set_mlir_data!(
         parent(x),
         get_mlir_data(@opcall(reshape(TracedRArray{T}(data), size(parent(x))...))),
@@ -116,7 +116,7 @@ function set_mlir_data!(x::Base.ReshapedArray{TracedRNumber{T}}, data) where {T}
 end
 
 function get_ancestor_and_indices(
-    x::Base.ReshapedArray{TracedRNumber{T},N}, indices::Vector{CartesianIndex{N}}
+    x::Base.ReshapedArray{<:TracedRNumber{T},N}, indices::Vector{CartesianIndex{N}}
 ) where {T,N}
     linear_indices = LinearIndices(size(x))[indices]
     parent_linear_indices = LinearIndices(size(parent(x)))[linear_indices]
@@ -124,7 +124,7 @@ function get_ancestor_and_indices(
 end
 
 function get_ancestor_and_indices(
-    x::Base.ReshapedArray{TracedRNumber{T},N}, indices...
+    x::Base.ReshapedArray{<:TracedRNumber{T},N}, indices...
 ) where {T,N}
     @assert length(indices) == N "Expected $N indices, got $(length(indices))"
     indices = Base.to_indices(x, indices)
@@ -160,7 +160,7 @@ function get_ancestor_and_indices(
 end
 
 function set_mlir_data!(
-    x::PermutedDimsArray{TracedRNumber{T},N,perm,iperm}, data
+    x::PermutedDimsArray{<:TracedRNumber{T},N,perm,iperm}, data
 ) where {T,N,perm,iperm}
     set_mlir_data!(parent(x), get_mlir_data(permutedims(TracedRArray{T}(data), iperm)))
     return x
diff --git a/src/stdlibs/LinearAlgebra.jl b/src/stdlibs/LinearAlgebra.jl
index 0bf5a9cd48..1c4b343763 100644
--- a/src/stdlibs/LinearAlgebra.jl
+++ b/src/stdlibs/LinearAlgebra.jl
@@ -81,7 +81,7 @@ include("factorization/Factorization.jl")
 
 # Various Wrapper Arrays defined in LinearAlgebra
 function ReactantCore.materialize_traced_array(
-    x::Transpose{TracedRNumber{T},<:AnyTracedRArray}
+    x::Transpose{<:TracedRNumber{T},<:AnyTracedRArray}
 ) where {T}
     px = materialize_traced_array(parent(x))
     A = ndims(px) == 1 ? reshape(px, :, 1) : px
@@ -89,7 +89,7 @@ function ReactantCore.materialize_traced_array(
 end
 
 function ReactantCore.materialize_traced_array(
-    x::Adjoint{TracedRNumber{T},<:AnyTracedRArray}
+    x::Adjoint{<:TracedRNumber{T},<:AnyTracedRArray}
 ) where {T}
     return @opcall conj(
         materialize_traced_array(transpose(materialize_traced_array(parent(x))))
@@ -97,13 +97,13 @@ function ReactantCore.materialize_traced_array(
 end
 
 function ReactantCore.materialize_traced_array(
-    x::Diagonal{TracedRNumber{T},<:AnyTracedRVector}
+    x::Diagonal{<:TracedRNumber{T},<:AnyTracedRVector}
 ) where {T}
     return diagm(materialize_traced_array(parent(x)))
 end
 
 function ReactantCore.materialize_traced_array(
-    x::Tridiagonal{TracedRNumber{T},<:AnyTracedRVector}
+    x::Tridiagonal{<:TracedRNumber{T},<:AnyTracedRVector}
 ) where {T}
     return diagm(-1 => x.dl, 0 => x.d, 1 => x.du)
 end
@@ -112,7 +112,7 @@ for (AT, comp) in ((:LowerTriangular, "GE"), (:UpperTriangular, "LE"))
     uAT = Symbol(:Unit, AT)
     @eval begin
         function ReactantCore.materialize_traced_array(
-            x::LinearAlgebra.$(AT){TracedRNumber{T},<:AnyTracedRMatrix}
+            x::LinearAlgebra.$(AT){<:TracedRNumber{T},<:AnyTracedRMatrix}
         ) where {T}
             m, n = size(x)
             px = materialize_traced_array(parent(x))
@@ -123,7 +123,7 @@ for (AT, comp) in ((:LowerTriangular, "GE"), (:UpperTriangular, "LE"))
         end
 
         function ReactantCore.materialize_traced_array(
-            x::LinearAlgebra.$(uAT){TracedRNumber{T},<:AnyTracedRMatrix}
+            x::LinearAlgebra.$(uAT){<:TracedRNumber{T},<:AnyTracedRMatrix}
         ) where {T}
             m, n = size(x)
             px = materialize_traced_array(parent(x))
@@ -139,7 +139,7 @@ for (AT, comp) in ((:LowerTriangular, "GE"), (:UpperTriangular, "LE"))
 end
 
 function ReactantCore.materialize_traced_array(
-    x::Hermitian{TracedRNumber{T},<:AnyTracedRMatrix}
+    x::Hermitian{<:TracedRNumber{T},<:AnyTracedRMatrix}
 ) where {T}
     m, n = size(x)
     row_idxs = @opcall iota(Int, [m, n]; iota_dimension=1)
@@ -152,7 +152,7 @@ function ReactantCore.materialize_traced_array(
 end
 
 function ReactantCore.materialize_traced_array(
-    x::Symmetric{TracedRNumber{T},<:AnyTracedRMatrix}
+    x::Symmetric{<:TracedRNumber{T},<:AnyTracedRMatrix}
 ) where {T}
     m, n = size(x)
     row_idxs = @opcall iota(Int, [m, n]; iota_dimension=1)
@@ -165,7 +165,7 @@ function ReactantCore.materialize_traced_array(
 end
 
 function TracedUtils.set_mlir_data!(
-    x::Transpose{TracedRNumber{T},TracedRArray{T,N}}, data
+    x::Transpose{<:TracedRNumber{T},<:TracedRArray{T,N}}, data
 ) where {T,N}
     tdata = TracedRArray{T}(data)
     px = parent(x)
@@ -180,7 +180,7 @@ function TracedUtils.set_mlir_data!(
 end
 
 function TracedUtils.set_mlir_data!(
-    x::Adjoint{TracedRNumber{T},TracedRArray{T,N}}, data
+    x::Adjoint{<:TracedRNumber{T},<:TracedRArray{T,N}}, data
 ) where {T,N}
     tdata = TracedRArray{T}(data)
     px = parent(x)
@@ -194,7 +194,7 @@ function TracedUtils.set_mlir_data!(
 end
 
 function TracedUtils.set_mlir_data!(
-    x::Diagonal{TracedRNumber{T},TracedRArray{T,1}}, data
+    x::Diagonal{<:TracedRNumber{T},<:TracedRArray{T,1}}, data
 ) where {T}
     parent(x).mlir_data = diag(TracedRArray{T}(data)).mlir_data
     return x
@@ -207,7 +207,7 @@ for (AT, dcomp, ocomp) in (
     (:UnitUpperTriangular, "LT", "GE"),
 )
     @eval function TracedUtils.set_mlir_data!(
-        x::LinearAlgebra.$(AT){TracedRNumber{T},<:AnyTracedRMatrix}, data
+        x::LinearAlgebra.$(AT){<:TracedRNumber{T},<:AnyTracedRMatrix}, data
     ) where {T}
         tdata = TracedRArray{T}(data)
         z = zero(tdata)
@@ -228,7 +228,7 @@ for (AT, dcomp, ocomp) in (
 end
 
 function TracedUtils.set_mlir_data!(
-    x::Hermitian{TracedRNumber{T},<:AnyTracedRMatrix}, data
+    x::Hermitian{<:TracedRNumber{T},<:AnyTracedRMatrix}, data
 ) where {T}
     if x.uplo == 'L'
         set_mlir_data!(LowerTriangular(parent(x)), data)
@@ -239,7 +239,7 @@ function TracedUtils.set_mlir_data!(
 end
 
 function TracedUtils.set_mlir_data!(
-    x::Symmetric{TracedRNumber{T},<:AnyTracedRMatrix}, data
+    x::Symmetric{<:TracedRNumber{T},<:AnyTracedRMatrix}, data
 ) where {T}
     if x.uplo == 'L'
         set_mlir_data!(LowerTriangular(parent(x)), data)
@@ -250,7 +250,7 @@ function TracedUtils.set_mlir_data!(
 end
 
 function TracedUtils.set_mlir_data!(
-    x::Tridiagonal{TracedRNumber{T},<:AnyTracedRVector}, data
+    x::Tridiagonal{<:TracedRNumber{T},<:AnyTracedRVector}, data
 ) where {T}
     tdata = TracedRArray{T}(data)
     set_mlir_data!(x.dl, materialize_traced_array(diag(tdata, -1)).mlir_data)
@@ -259,7 +259,7 @@ function TracedUtils.set_mlir_data!(
     return x
 end
 
-Reactant.aos_to_soa(x::Tridiagonal{TracedRNumber{T}}) where {T} = x
+Reactant.aos_to_soa(x::Tridiagonal{<:TracedRNumber{T}}) where {T} = x
 
 # Core functions
 function overloaded_mul(
@@ -622,7 +622,7 @@ tfun_to_char(::typeof(transpose)) = 'T'
 tfun_to_char(::typeof(adjoint)) = 'C'
 
 function LinearAlgebra.generic_trimatdiv!(
-    C::AbstractVecOrMat{TracedRNumber{T}},
+    C::AbstractVecOrMat{<:TracedRNumber{T}},
     uploc,
     isunitc,
     tfun::Function,
@@ -645,7 +645,7 @@ function LinearAlgebra.generic_trimatdiv!(
 end
 
 function LinearAlgebra.generic_trimatdiv!(
-    C::AbstractVecOrMat{TracedRNumber{T}},
+    C::AbstractVecOrMat{<:TracedRNumber{T}},
     uploc,
     isunitc,
     tfun::Function,
@@ -659,7 +659,7 @@ function LinearAlgebra.generic_trimatdiv!(
 end
 
 function LinearAlgebra.generic_mattridiv!(
-    C::AbstractMatrix{TracedRNumber{T}},
+    C::AbstractMatrix{<:TracedRNumber{T}},
     uploc,
     isunitc,
     tfun::Function,
@@ -682,7 +682,7 @@ function LinearAlgebra.generic_mattridiv!(
 end
 
 function LinearAlgebra.generic_mattridiv!(
-    C::AbstractMatrix{TracedRNumber{T}},
+    C::AbstractMatrix{<:TracedRNumber{T}},
     uploc,
     isunitc,
     tfun::Function,

From 25f12ff4cc6c1b620744b0d95d2b3ecdf7d87a5d Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Wed, 8 Jul 2026 23:03:07 +0200
Subject: [PATCH 02/30] Make `similar` of a concrete array with a traced eltype
 return a traced 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.
---
 src/ConcreteRArray.jl              | 14 ++++++++++++++
 test/integration/linear_algebra.jl | 13 +++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index c807ba9ab0..cc341ded46 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -484,6 +484,15 @@ end
 
 Base.similar(a::ConcretePJRTArray, dims::Dims) = similar(a, eltype(a), dims)
 
+# A concrete array asked for a traced element type (e.g. via LinearAlgebra's
+# generic `similar`/`\` while tracing) cannot hold traced values in its device
+# buffer, so produce a traced array instead.
+function Base.similar(
+    ::ConcretePJRTArray{T,N,D}, ::Type{S}, dims::Dims
+) where {S<:TracedRNumber,T,N,D}
+    return similar(TracedRArray{unwrapped_eltype(S)}, dims)
+end
+
 function Base.similar(AT::Type{<:ConcretePJRTArray{T}}, dims::Dims; kwargs...) where {T}
     return similar(AT, T, dims; kwargs...)
 end
@@ -493,6 +502,11 @@ function Base.similar(a::ConcreteIFRTArray{T}, ::Type{S}=T, dims::Dims=size(a))
         Array{S}(undef, dims); client=XLA.client(a), device=XLA.device(a), a.sharding
     )
 end
+function Base.similar(
+    ::ConcreteIFRTArray{T}, ::Type{S}, dims::Dims
+) where {S<:TracedRNumber,T}
+    return similar(TracedRArray{unwrapped_eltype(S)}, dims)
+end
 Base.similar(a::ConcreteIFRTArray, dims::Dims) = similar(a, eltype(a), dims)
 function Base.similar(::Type{ConcreteIFRTArray{T}}, dims::Dims) where {T}
     return ConcreteIFRTArray{T}(undef, dims)
diff --git a/test/integration/linear_algebra.jl b/test/integration/linear_algebra.jl
index c301db6d8e..01286cc820 100644
--- a/test/integration/linear_algebra.jl
+++ b/test/integration/linear_algebra.jl
@@ -546,3 +546,16 @@ diagonal_3_arg_mul(x, y) = y' * Diagonal(x) * y
 
     @test @jit(diagonal_3_arg_mul(x_ra, y_ra)) ≈ diagonal_3_arg_mul(x, y)
 end
+
+@testset "Triangular solve with captured concrete RHS" begin
+    M = collect(reshape(1.0:9.0, 3, 3)) + 10.0 * I
+    b = [1.0, 2.0, 3.0]
+    M_ra = Reactant.to_rarray(M)
+    b_ra = Reactant.to_rarray(b)
+    # `b_ra` is captured (concrete); the solve allocates its output via
+    # `similar(b_ra, <traced eltype>, ...)`, which must yield a traced array
+    # rather than hit `primitive_type` on the wrapped element type.
+    @test Array(@jit((m -> UpperTriangular(m) \ b_ra)(M_ra))) ≈ UpperTriangular(M) \ b
+    @test Array(@jit(((m, v) -> LowerTriangular(m) \ v)(M_ra, b_ra))) ≈
+        LowerTriangular(M) \ b
+end

From 74e69bb10dfd77ac74f65007b4222041b9c486bd Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Tue, 7 Jul 2026 20:13:16 +0200
Subject: [PATCH 03/30] Split Reactant number types by numeric kind

Traced and concrete Reactant numbers are now split into integer, float,
and complex types so that they subtype Integer and AbstractFloat where
possible (issue #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.
---
 src/Reactant.jl |  19 +++-
 src/Types.jl    | 233 ++++++++++++++++++++++++++++++++++--------------
 2 files changed, 179 insertions(+), 73 deletions(-)

diff --git a/src/Reactant.jl b/src/Reactant.jl
index 789f352b27..f02978433b 100644
--- a/src/Reactant.jl
+++ b/src/Reactant.jl
@@ -123,15 +123,17 @@ isa_traced_soa(::AbstractRange{<:TracedRNumber}) = true
 
 unwrapped_eltype(::Type{T}) where {T<:Number} = T
 unwrapped_eltype(::Type{<:RNumber{T}}) where {T} = T
-unwrapped_eltype(::Type{TracedRNumber{T}}) where {T} = T
 
 unwrapped_eltype(::T) where {T<:Number} = T
 unwrapped_eltype(::RNumber{T}) where {T} = T
-unwrapped_eltype(::TracedRNumber{T}) where {T} = T
 
 unwrapped_eltype(::Type{<:AbstractArray{T}}) where {T} = unwrapped_eltype(T)
 unwrapped_eltype(::AbstractArray{T}) where {T} = unwrapped_eltype(T)
 
+# `TracedRArray{T,N}` with the element-type parameter left free does not match
+# `Type{<:AbstractArray{T}}` (no single eltype), so extract `T` directly.
+unwrapped_eltype(::Type{<:TracedRArray{T}}) where {T} = T
+
 include("Ops.jl")
 Base.push!(no_rewrite_ancestor_modules, Ops)
 
@@ -157,7 +159,7 @@ aos_to_soa(x::AbstractArray) = x
 aos_to_soa(x::TracedRArray) = x
 aos_to_soa(x::AnyTracedRArray) = x
 
-function aos_to_soa(x::Array{TracedRNumber{T}}) where {T}
+function aos_to_soa(x::Array{<:TracedRNumber{T}}) where {T}
     isa_traced_soa(ancestor(x)) && return x
     for i in eachindex(x)
         if !isassigned(x, i)
@@ -293,12 +295,21 @@ using .Compiler:
     code_xla,
     traced_getfield,
     compile
-export ConcreteRArray,
+export RInteger,
+    RFloat,
+    RComplex,
+    ConcreteRArray,
     ConcreteRNumber,
     ConcretePJRTArray,
     ConcretePJRTNumber,
+    ConcretePJRTInteger,
+    ConcretePJRTFloat,
+    ConcretePJRTComplex,
     ConcreteIFRTArray,
     ConcreteIFRTNumber,
+    ConcreteIFRTInteger,
+    ConcreteIFRTFloat,
+    ConcreteIFRTComplex,
     @compile,
     @code_hlo,
     @code_mhlo,
diff --git a/src/Types.jl b/src/Types.jl
index 528ff17600..0d34dec85a 100644
--- a/src/Types.jl
+++ b/src/Types.jl
@@ -1,6 +1,19 @@
-abstract type RNumber{T<:ReactantPrimitive} <: Number end
+# Reactant numbers are split by numeric kind so that traced and concrete values
+# subtype `Integer` and `AbstractFloat` where possible. Julia has no abstract
+# complex type, so complex Reactant numbers can only subtype `Number`.
+abstract type RInteger{T<:ReactantPrimitive} <: Integer end
+abstract type RFloat{T<:ReactantPrimitive} <: AbstractFloat end
+abstract type RComplex{T<:ReactantPrimitive} <: Number end
 
-abstract type AbstractConcreteNumber{T} <: RNumber{T} end
+const RReal{T} = Union{RInteger{T},RFloat{T}}
+const RNumber{T} = Union{RReal{T},RComplex{T}}
+
+abstract type AbstractConcreteInteger{T} <: RInteger{T} end
+abstract type AbstractConcreteFloat{T} <: RFloat{T} end
+abstract type AbstractConcreteComplex{T} <: RComplex{T} end
+
+const AbstractConcreteReal{T} = Union{AbstractConcreteInteger{T},AbstractConcreteFloat{T}}
+const AbstractConcreteNumber{T} = Union{AbstractConcreteReal{T},AbstractConcreteComplex{T}}
 
 abstract type RArray{T,N} <: DenseArray{T,N} end
 
@@ -37,63 +50,94 @@ end
 ## MissingTracedValue -- defined in ReactantCore
 @leaf MissingTracedValue
 
-## TracedRNumber
-mutable struct TracedRNumber{T} <: RNumber{T}
-    paths::Tuple
-    mlir_data::Union{Nothing,MLIR.IR.Value}
-
-    function TracedRNumber{T}(
-        paths::Tuple, mlir_data::Union{Nothing,MLIR.IR.Value}
-    ) where {T}
-        if !isnothing(mlir_data)
-            @assert size(MLIR.IR.type(mlir_data)) == ()
+## TracedRInteger, TracedRFloat, TracedRComplex
+for (TracedT, AbstractT) in
+    ((:TracedRInteger, :RInteger), (:TracedRFloat, :RFloat), (:TracedRComplex, :RComplex))
+    @eval begin
+        mutable struct $TracedT{T} <: $AbstractT{T}
+            paths::Tuple
+            mlir_data::Union{Nothing,MLIR.IR.Value}
+
+            function $TracedT{T}(
+                paths::Tuple, mlir_data::Union{Nothing,MLIR.IR.Value}
+            ) where {T}
+                if !isnothing(mlir_data)
+                    @assert size(MLIR.IR.type(mlir_data)) == ()
+                end
+                return new{T}(paths, mlir_data)
+            end
         end
-        return new{T}(paths, mlir_data)
+
+        @leaf $TracedT
     end
 end
 
-Base.elsize(::Type{TracedRNumber{T}}) where {T} = sizeof(T)
-Base.elsize(::Type{RNumber{T}}) where {T} = sizeof(T)
-Base.elsize(::Type{<:AbstractConcreteNumber{T}}) where {T} = sizeof(T)
-Base.elsize(::Type{<:AbstractConcreteArray{T}}) where {T} = sizeof(T)
+const TracedRReal{T} = Union{TracedRInteger{T},TracedRFloat{T}}
+const TracedRNumber{T} = Union{TracedRReal{T},TracedRComplex{T}}
+
+"""
+    traced_number_type(::Type{T}) where {T<:ReactantPrimitive}
 
-function repath(x::TracedRNumber{T}, paths) where {T}
-    return TracedRNumber{T}(paths, x.mlir_data)
+The concrete traced number type for element type `T`, i.e. the member of
+`TracedRNumber{T}` that matches the numeric kind of `T`.
+"""
+@inline traced_number_type(::Type{T}) where {T<:Integer} = TracedRInteger{T}
+@inline traced_number_type(::Type{T}) where {T<:AbstractFloat} = TracedRFloat{T}
+@inline traced_number_type(::Type{T}) where {T<:Complex} = TracedRComplex{T}
+
+@inline function TracedRNumber{T}(
+    paths::Tuple, mlir_data::Union{Nothing,MLIR.IR.Value}
+) where {T}
+    return traced_number_type(T)(paths, mlir_data)
 end
 
-@leaf TracedRNumber
+Base.elsize(::Type{<:RNumber{T}}) where {T} = sizeof(T)
+Base.elsize(::Type{<:AbstractConcreteArray{T}}) where {T} = sizeof(T)
+
+repath(x::TracedRNumber, paths) = Core.Typeof(x)(paths, x.mlir_data)
 
 ## TracedRArray
-mutable struct TracedRArray{T,N} <: RArray{TracedRNumber{T},N}
+## The element type `RT` is redundant (it is always `traced_number_type(T)`) but
+## exposes the numeric kind of the elements in the `RArray` supertype, so that
+## e.g. a traced float array is an `AbstractArray{<:AbstractFloat}`.
+mutable struct TracedRArray{T,N,RT<:TracedRNumber{T}} <: RArray{RT,N}
     paths::Tuple
     mlir_data::Union{Nothing,MLIR.IR.Value}
     shape::NTuple{N,Int}
 
-    function TracedRArray{T,N}(
+    function TracedRArray{T,N,RT}(
         paths::Tuple, mlir_data::Union{Nothing,MLIR.IR.Value}, shape
-    ) where {T,N}
+    ) where {T,N,RT<:TracedRNumber{T}}
         shape = Tuple(shape)
         if !isnothing(mlir_data)
             @assert size(MLIR.IR.type(mlir_data)) == shape "Expected: $(shape), got: $(size(MLIR.IR.type(mlir_data)))"
         end
-        return new{T,N}(paths, mlir_data, shape)
+        return new{T,N,RT}(paths, mlir_data, shape)
     end
+end
 
-    function TracedRArray{T,N}(::UndefInitializer, shape::Integer...) where {T,N}
-        return similar(TracedRArray{T,N}, shape...)
-    end
+function TracedRArray{T,N}(
+    paths::Tuple, mlir_data::Union{Nothing,MLIR.IR.Value}, shape
+) where {T,N}
+    return TracedRArray{T,N,traced_number_type(T)}(paths, mlir_data, shape)
+end
 
-    function TracedRArray{T,N}(::UndefInitializer, shape::NTuple{N,Int}) where {T,N}
-        return similar(TracedRArray{T,N}, shape)
-    end
+function TracedRArray{T,N}(::UndefInitializer, shape::Integer...) where {T,N}
+    return similar(TracedRArray{T,N}, shape...)
+end
+
+function TracedRArray{T,N}(::UndefInitializer, shape::NTuple{N,Int}) where {T,N}
+    return similar(TracedRArray{T,N}, shape)
 end
 
 function repath(x::TracedRArray{T,N}, paths) where {T,N}
-    return TracedRArray{T,N}(paths, x.mlir_data, x.shape)
+    return Core.Typeof(x)(paths, x.mlir_data, x.shape)
 end
 
 @leaf TracedRArray
-Adapt.parent_type(::Type{TracedRArray{T,N}}) where {T,N} = TracedRArray{T,N}
+function Adapt.parent_type(::Type{<:TracedRArray{T,N}}) where {T,N}
+    return TracedRArray{T,N,traced_number_type(T)}
+end
 
 ## TracedStepRangeLen
 struct TracedStepRangeLen{T,R,S,L} <: AbstractRange{T}
@@ -135,9 +179,7 @@ end
 Adapt.parent_type(::Type{TracedUnitRange{T}}) where {T} = TracedUnitRange{T}
 
 ## TracedRational
-struct TracedRational{
-    T<:Union{<:Integer,<:AbstractConcreteNumber{<:Integer},TracedRNumber{<:Integer}}
-} <: Real
+struct TracedRational{T<:Integer} <: Real
     num::T
     den::T
 end
@@ -151,21 +193,51 @@ const AnyTracedRMatrix{T} = AnyTracedRArray{T,2}
 const AnyTracedRVecOrMat{T} = Union{AnyTracedRVector{T},AnyTracedRMatrix{T}}
 
 # Concrete Types
-## ConcretePJRTNumber
-mutable struct ConcretePJRTNumber{T,D} <: AbstractConcreteNumber{T}
-    data::NTuple{D,XLA.PJRT.AsyncBuffer}
-    sharding::Sharding.ShardInfo
-    donated::Bool
+## ConcretePJRTInteger, ConcretePJRTFloat, ConcretePJRTComplex
+for (ConcreteT, AbstractT) in (
+    (:ConcretePJRTInteger, :AbstractConcreteInteger),
+    (:ConcretePJRTFloat, :AbstractConcreteFloat),
+    (:ConcretePJRTComplex, :AbstractConcreteComplex),
+)
+    @eval begin
+        mutable struct $ConcreteT{T,D} <: $AbstractT{T}
+            data::NTuple{D,XLA.PJRT.AsyncBuffer}
+            sharding::Sharding.ShardInfo
+            donated::Bool
+
+            function $ConcreteT{T,D}(
+                data::NTuple{D,XLA.PJRT.AsyncBuffer}, sharding::Sharding.ShardInfo
+            ) where {T,D}
+                return new{T,D}(data, sharding, false)
+            end
+        end
 
-    function ConcretePJRTNumber{T,D}(
-        data::NTuple{D,XLA.PJRT.AsyncBuffer}, sharding::Sharding.ShardInfo
-    ) where {T,D}
-        return new{T,D}(data, sharding, false)
+        @leaf $ConcreteT
     end
 end
 
+const ConcretePJRTReal{T,D} = Union{ConcretePJRTInteger{T,D},ConcretePJRTFloat{T,D}}
+const ConcretePJRTNumber{T,D} = Union{ConcretePJRTReal{T,D},ConcretePJRTComplex{T,D}}
+
+"""
+    pjrt_number_type(::Type{T}) where {T<:ReactantPrimitive}
+
+The concrete PJRT number type for element type `T` (with the device-count
+parameter `D` left free), i.e. the member of `ConcretePJRTNumber{T}` that
+matches the numeric kind of `T`.
+"""
+@inline pjrt_number_type(::Type{T}) where {T<:Integer} = ConcretePJRTInteger{T}
+@inline pjrt_number_type(::Type{T}) where {T<:AbstractFloat} = ConcretePJRTFloat{T}
+@inline pjrt_number_type(::Type{T}) where {T<:Complex} = ConcretePJRTComplex{T}
+
 ConcretePJRTNumber{T,1}(x::Number) where {T} = ConcretePJRTNumber{T}(x)
 
+function ConcretePJRTNumber{T,D}(
+    data::NTuple{D,XLA.PJRT.AsyncBuffer}, sharding::Sharding.ShardInfo
+) where {T,D}
+    return pjrt_number_type(T){D}(data, sharding)
+end
+
 function ConcretePJRTNumber{T}(data::Tuple{XLA.PJRT.AsyncBuffer}) where {T}
     return ConcretePJRTNumber{T,1}(data, Sharding.NoShardInfo())
 end
@@ -174,8 +246,6 @@ function ConcretePJRTNumber{T}(data::NTuple{D,XLA.PJRT.AsyncBuffer}, sharding) w
     return ConcretePJRTNumber{T,D}(data, sharding)
 end
 
-@leaf ConcretePJRTNumber
-
 function ConcretePJRTNumber{T}(data::T2; kwargs...) where {T<:Number,T2<:Number}
     carray = ConcretePJRTArray(fill(convert(T, data)); kwargs...)
     if !Sharding.is_sharded(carray.sharding)
@@ -296,25 +366,52 @@ end
 
 # While sharding is part of IFRT.Array, we still need to carry it around for compiling the
 # MLIR module.
-## ConcreteIFRTNumber
-mutable struct ConcreteIFRTNumber{T} <: AbstractConcreteNumber{T}
-    data::XLA.IFRT.AsyncArray
-    sharding::Sharding.ShardInfo
-    donated::Bool
+## ConcreteIFRTInteger, ConcreteIFRTFloat, ConcreteIFRTComplex
+for (ConcreteT, AbstractT) in (
+    (:ConcreteIFRTInteger, :AbstractConcreteInteger),
+    (:ConcreteIFRTFloat, :AbstractConcreteFloat),
+    (:ConcreteIFRTComplex, :AbstractConcreteComplex),
+)
+    @eval begin
+        mutable struct $ConcreteT{T} <: $AbstractT{T}
+            data::XLA.IFRT.AsyncArray
+            sharding::Sharding.ShardInfo
+            donated::Bool
+
+            function $ConcreteT{T}(
+                data::XLA.IFRT.AsyncArray, sharding::Sharding.ShardInfo
+            ) where {T}
+                return new{T}(data, sharding, false)
+            end
+        end
 
-    function ConcreteIFRTNumber{T}(
-        data::XLA.IFRT.AsyncArray, sharding::Sharding.ShardInfo
-    ) where {T}
-        return new{T}(data, sharding, false)
+        @leaf $ConcreteT
     end
 end
 
+const ConcreteIFRTReal{T} = Union{ConcreteIFRTInteger{T},ConcreteIFRTFloat{T}}
+const ConcreteIFRTNumber{T} = Union{ConcreteIFRTReal{T},ConcreteIFRTComplex{T}}
+
+"""
+    ifrt_number_type(::Type{T}) where {T<:ReactantPrimitive}
+
+The concrete IFRT number type for element type `T`, i.e. the member of
+`ConcreteIFRTNumber{T}` that matches the numeric kind of `T`.
+"""
+@inline ifrt_number_type(::Type{T}) where {T<:Integer} = ConcreteIFRTInteger{T}
+@inline ifrt_number_type(::Type{T}) where {T<:AbstractFloat} = ConcreteIFRTFloat{T}
+@inline ifrt_number_type(::Type{T}) where {T<:Complex} = ConcreteIFRTComplex{T}
+
+function ConcreteIFRTNumber{T}(
+    data::XLA.IFRT.AsyncArray, sharding::Sharding.ShardInfo
+) where {T}
+    return ifrt_number_type(T)(data, sharding)
+end
+
 function ConcreteIFRTNumber{T}(data::XLA.IFRT.AsyncArray) where {T}
     return ConcreteIFRTNumber{T}(data, Sharding.NoShardInfo())
 end
 
-@leaf ConcreteIFRTNumber
-
 function ConcreteIFRTNumber{T}(data::T2; kwargs...) where {T<:Number,T2<:Number}
     carray = ConcreteIFRTArray(fill(convert(T, data)); kwargs...)
     return ConcreteIFRTNumber{T}(carray.data, carray.sharding)
@@ -477,7 +574,7 @@ function InterpolateArray(
                 N_dim, M_dim = final_grid_size[dim], src_size[dim]
                 H = halo[dim]
 
-                [
+                return [
                     begin
                         if I <= H
                             clamp(I, 1, M_dim)
@@ -516,7 +613,7 @@ function InterpolateArray(
                 N_dim, M_dim = final_grid_size[dim], src_size[dim]
                 H = halo[dim]
 
-                [
+                return [
                     begin
                         if I <= H
                             clamp(I, 1, M_dim)
@@ -541,7 +638,7 @@ function InterpolateArray(
                 N_dim, M_dim = final_grid_size[dim], src_size[dim]
                 H = halo[dim]
 
-                [
+                return [
                     begin
                         if I <= H
                             clamp(I, 1, M_dim)
@@ -563,7 +660,7 @@ function InterpolateArray(
 
             dens = ntuple(N) do dim
                 H = halo[dim]
-                2 * max(1, final_grid_size[dim] - 2 * H)
+                return 2 * max(1, final_grid_size[dim] - 2 * H)
             end
             total_den = prod(dens)
 
@@ -572,7 +669,7 @@ function InterpolateArray(
                 N_dim, M_dim = final_grid_size[dim], src_size[dim]
                 H = halo[dim]
 
-                [
+                return [
                     begin
                         if I <= H || I >= N_dim - H + 1
                             0
@@ -597,13 +694,11 @@ function InterpolateArray(
                         dim -> c[dim] == 1 ? lows[dim][I.I[dim]] : highs[dim][I.I[dim]], N
                     )
 
-                    w_int = prod(
-                        ntuple(dim -> if c[dim] == 1
-                            (dens[dim] - rems[dim][I.I[dim]])
-                        else
-                            rems[dim][I.I[dim]]
-                        end, N)
-                    )
+                    w_int = prod(ntuple(dim -> if c[dim] == 1
+                        (dens[dim] - rems[dim][I.I[dim]])
+                    else
+                        rems[dim][I.I[dim]]
+                    end, N))
 
                     sum_val += w_int * local_cpu_array[CartesianIndex(idx)]
                 end

From 238936f5b449d09e13fbc6d1865423b3362604fb Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Tue, 7 Jul 2026 20:13:16 +0200
Subject: [PATCH 04/30] Adapt traced number operations to the split number
 types

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.
---
 src/TracedRNumber.jl | 534 ++++++++++++++++++++++---------------------
 1 file changed, 269 insertions(+), 265 deletions(-)

diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index 8289603137..9d2ed82601 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -1,6 +1,16 @@
 module TracedRNumberOverrides
 
-using ..Reactant: Reactant, TracedRNumber, TracedRArray, Ops, unwrapped_eltype
+using ..Reactant:
+    Reactant,
+    TracedRNumber,
+    TracedRReal,
+    TracedRInteger,
+    TracedRFloat,
+    TracedRComplex,
+    TracedRArray,
+    Ops,
+    unwrapped_eltype,
+    traced_number_type
 using ..Ops: @opcall
 using ReactantCore: ReactantCore, @trace
 using Adapt: Adapt
@@ -17,7 +27,8 @@ ReactantCore.is_traced(::TracedRNumber) = true
 
 Base.to_index(x::TracedRNumber{<:Integer}) = x
 
-Base.precision(x::TracedRNumber{T}; kwargs...) where {T} = precision(T; kwargs...)
+Base.precision(x::TracedRFloat{T}; kwargs...) where {T} = precision(T; kwargs...)
+Base.precision(::Type{<:TracedRFloat{T}}; kwargs...) where {T} = precision(T; kwargs...)
 
 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))
@@ -25,92 +36,96 @@ Base.collect(x::TracedRNumber{T}) where {T} = TracedRArray{T,0}((), x.mlir_data,
 
 Base.copy(x::TracedRNumber{T}) where {T} = TracedRNumber{T}((), x.mlir_data)
 
-function Base.eps(::Type{TracedRNumber{T}}) where {T}
-    return Reactant.promote_to(TracedRNumber{T}, eps(T))
+function Base.eps(::Type{<:TracedRFloat{T}}) where {T}
+    return Reactant.promote_to(TracedRFloat{T}, eps(T))
 end
-Base.eps(x::TracedRNumber{T}) where {T} = eps(typeof(x))
+Base.eps(x::TracedRFloat) = eps(typeof(x))
 
 function Base.typemin(::Type{<:TracedRNumber{T}}) where {T}
     return Reactant.promote_to(TracedRNumber{T}, typemin(T))
 end
-Base.typemin(x::TracedRNumber{T}) where {T} = typemin(typeof(x))
+Base.typemin(x::TracedRNumber) = typemin(typeof(x))
 
 function Base.typemax(::Type{<:TracedRNumber{T}}) where {T}
     return Reactant.promote_to(TracedRNumber{T}, typemax(T))
 end
-Base.typemax(x::TracedRNumber{T}) where {T} = typemax(typeof(x))
+Base.typemax(x::TracedRNumber) = typemax(typeof(x))
 
-Base.floatmin(::Type{Reactant.TracedRNumber{T}}) where {T} = floatmin(T)
-Base.floatmax(::Type{Reactant.TracedRNumber{T}}) where {T} = floatmax(T)
+Base.floatmin(::Type{<:TracedRFloat{T}}) where {T} = floatmin(T)
+Base.floatmax(::Type{<:TracedRFloat{T}}) where {T} = floatmax(T)
 
-function Base.nextfloat(x::TracedRNumber{T}) where {T<:AbstractFloat}
-    return @opcall next_after(x, typemax(x))
-end
-
-function Base.prevfloat(x::TracedRNumber{T}) where {T<:AbstractFloat}
-    return @opcall next_after(x, typemin(x))
-end
+Base.nextfloat(x::TracedRFloat) = @opcall next_after(x, typemax(x))
+Base.prevfloat(x::TracedRFloat) = @opcall next_after(x, typemin(x))
 
 function Base.rtoldefault(T::Type{<:TracedRNumber})
     return T(Base.rtoldefault(unwrapped_eltype(T)))
 end
+function Base.rtoldefault(T::Type{<:TracedRFloat})
+    return T(Base.rtoldefault(unwrapped_eltype(T)))
+end
 
-function Base.isfinite(x::TracedRNumber{<:Complex})
+function Base.isfinite(x::TracedRComplex)
     return isfinite(real(x)) & isfinite(imag(x))
 end
-Base.isfinite(x::TracedRNumber{<:AbstractFloat}) = @opcall is_finite(x)
+Base.isfinite(x::TracedRFloat) = @opcall is_finite(x)
 
-function Base.isnan(x::TracedRNumber{<:Complex})
+function Base.isnan(x::TracedRComplex)
     return isnan(real(x)) | isnan(imag(x))
 end
-function Base.isnan(x::TracedRNumber{T}) where {T<:AbstractFloat}
+function Base.isnan(x::TracedRFloat{T}) where {T}
     return !isfinite(x) & (x != typemax(T)) & (x != typemin(T))
 end
 
-Base.isinf(x::TracedRNumber{<:Complex}) = isinf(real(x)) | isinf(imag(x))
-Base.isinf(x::TracedRNumber{<:AbstractFloat}) = @opcall is_inf(x)
-Base.isinf(::TracedRNumber{<:Integer}) = false
+Base.isinf(x::TracedRComplex) = isinf(real(x)) | isinf(imag(x))
+Base.isinf(x::TracedRFloat) = @opcall is_inf(x)
 
-function Base.show(io::IOty, X::TracedRNumber{T}) where {T,IOty<:Union{IO,IOContext}}
-    return print(io, "TracedRNumber{", T, "}(", X.paths, ")")
+for TracedT in (:TracedRInteger, :TracedRFloat, :TracedRComplex)
+    @eval function Base.show(io::IO, X::$TracedT{T}) where {T}
+        return print(io, $(string(TracedT)), "{", T, "}(", X.paths, ")")
+    end
 end
 
 Base.only(A::TracedRNumber{T}) where {T} = A
 
-function Base.promote_rule(::Type{TracedRNumber{T}}, ::Type{TracedRNumber{S}}) where {T,S}
-    return TracedRNumber{Base.promote_type(T, S)}
+# The result of promoting the primitive types selects the traced leaf type, so
+# a single parametric rule covers integer/float/complex crossings.
+function Base.promote_rule(
+    ::Type{<:TracedRNumber{T}}, ::Type{<:TracedRNumber{S}}
+) where {T,S}
+    return traced_number_type(Base.promote_type(T, S))
 end
 
-# Bool has special promotion rules in Base
-function Base.promote_rule(::Type{Bool}, ::Type{TracedRNumber{T}}) where {T}
-    return TracedRNumber{T}
+function Base.promote_rule(::Type{S}, ::Type{<:TracedRNumber{T}}) where {T,S<:Number}
+    return traced_number_type(Base.promote_type(T, S))
 end
 
-function Base.promote_rule(::Type{TracedRNumber{T}}, ::Type{Bool}) where {T}
-    return TracedRNumber{T}
+function Base.promote_rule(::Type{<:TracedRNumber{T}}, ::Type{S}) where {T,S<:Number}
+    return traced_number_type(Base.promote_type(T, S))
 end
 
-function Base.promote_rule(::Type{T}, ::Type{TracedRNumber{S}}) where {T,S}
-    return TracedRNumber{Base.promote_type(T, S)}
+# Bool has special promotion rules in Base
+function Base.promote_rule(::Type{Bool}, ::Type{<:TracedRNumber{T}}) where {T}
+    return traced_number_type(T)
 end
-
-function Base.promote_rule(::Type{TracedRNumber{T}}, ::Type{S}) where {T,S}
-    return TracedRNumber{Base.promote_type(T, S)}
+function Base.promote_rule(::Type{<:TracedRNumber{T}}, ::Type{Bool}) where {T}
+    return traced_number_type(T)
 end
 
+# Disambiguate against Base's irrational-vs-Real rule; real-only since
+# `TracedRReal <: Real` while the full union is not.
 function Base.promote_rule(
-    T::Type{<:AbstractIrrational}, ::Type{Reactant.TracedRNumber{S}}
+    T::Type{<:AbstractIrrational}, ::Type{<:TracedRReal{S}}
 ) where {S}
-    return TracedRNumber{Base.promote_type(T, S)}
+    return traced_number_type(Base.promote_type(T, S))
 end
 
 function Base.promote_rule(
-    ::Type{Reactant.TracedRNumber{S}}, T::Type{<:AbstractIrrational}
+    ::Type{<:TracedRReal{S}}, T::Type{<:AbstractIrrational}
 ) where {S}
-    return TracedRNumber{Base.promote_type(T, S)}
+    return traced_number_type(Base.promote_type(T, S))
 end
 
-function Base.promote_rule(::Type{Nothing}, ::Type{TracedRNumber{S}}) where {S}
+function Base.promote_rule(::Type{Nothing}, ::Type{<:TracedRNumber{S}}) where {S}
     return Union{Nothing,TracedRNumber{S}}
 end
 
@@ -168,11 +183,15 @@ function Base.promote_rule(::Type{TracedRNumber{T}}, S::Type{>:Nothing}) where {
     return Union{Nothing,TracedRNumber{Base.nonnothingtype(promote_type(T, S))}}
 end
 
-function Base.promote_rule(::Type{TwicePrecision{T}}, ::Type{TracedRNumber{S}}) where {T,S}
+function Base.promote_rule(
+    ::Type{TwicePrecision{T}}, ::Type{<:TracedRNumber{S}}
+) where {T,S}
     return TwicePrecision{Base.promote_type(T, TracedRNumber{S})}
 end
 
-function Base.promote_rule(::Type{TracedRNumber{T}}, ::Type{TwicePrecision{S}}) where {T,S}
+function Base.promote_rule(
+    ::Type{<:TracedRNumber{T}}, ::Type{TwicePrecision{S}}
+) where {T,S}
     return TwicePrecision{Base.promote_type(TracedRNumber{T}, S)}
 end
 
@@ -186,15 +205,14 @@ function TracedRNumber{T}(x::Number) where {T}
     return Reactant.promote_to(TracedRNumber{unwrapped_eltype(T)}, x)
 end
 
+# Covers the concrete traced types, e.g. via `convert(TracedRFloat{Float64}, x)`.
+(::Type{RT})(x::Number) where {RT<:TracedRNumber} = Reactant.promote_to(RT, x)
+
 for T in Base.uniontypes(Reactant.ReactantFloat8)
     @eval TracedRNumber{T}(x::$T) where {T} = Reactant.promote_to(TracedRNumber{T}, x)
 end
 
-for (aT, bT) in (
-    (TracedRNumber{<:Real}, Real),
-    (Real, TracedRNumber{<:Real}),
-    (TracedRNumber{<:Real}, TracedRNumber{<:Real}),
-)
+for (aT, bT) in ((TracedRReal, Real), (Real, TracedRReal), (TracedRReal, TracedRReal))
     @eval function Base.Complex(a::$aT, b::$bT)
         T = promote_type(unwrapped_eltype(a), unwrapped_eltype(b))
         a = Reactant.promote_to(TracedRNumber{T}, a)
@@ -203,34 +221,36 @@ for (aT, bT) in (
     end
 end
 
-Base.Complex(x::TracedRNumber{<:Real}) = @opcall complex(x, zero(x))
-Base.Complex(x::TracedRNumber{<:Complex}) = x
+Base.Complex(x::TracedRReal) = @opcall complex(x, zero(x))
+Base.Complex(x::TracedRComplex) = x
 
 # Base.complex
-Base.complex(::Type{TracedRNumber{T}}) where {T} = TracedRNumber{complex(T)}
-Base.complex(x::TracedRNumber{<:Real}) = complex(x, zero(x))
-function Base.complex(x::TracedRNumber{<:Real}, y::TracedRNumber{<:Real})
+Base.complex(::Type{TracedRNumber{T}}) where {T} = traced_number_type(complex(T))
+Base.complex(::Type{<:TracedRReal{T}}) where {T} = traced_number_type(complex(T))
+Base.complex(::Type{<:TracedRComplex{T}}) where {T} = TracedRComplex{T}
+Base.complex(x::TracedRReal) = complex(x, zero(x))
+function Base.complex(x::TracedRReal, y::TracedRReal)
     T = promote_type(unwrapped_eltype(x), unwrapped_eltype(y))
     return complex(
         Reactant.promote_to(TracedRNumber{T}, x), Reactant.promote_to(TracedRNumber{T}, y)
     )
 end
-function Base.complex(x::TracedRNumber{<:Real}, y::Real)
+function Base.complex(x::TracedRReal, y::Real)
     T = promote_type(unwrapped_eltype(x), typeof(y))
     return complex(
         Reactant.promote_to(TracedRNumber{T}, x), Reactant.promote_to(TracedRNumber{T}, y)
     )
 end
-function Base.complex(x::Real, y::TracedRNumber{<:Real})
+function Base.complex(x::Real, y::TracedRReal)
     T = promote_type(typeof(x), unwrapped_eltype(y))
     return complex(
         Reactant.promote_to(TracedRNumber{T}, x), Reactant.promote_to(TracedRNumber{T}, y)
     )
 end
-function Base.complex(x::TracedRNumber{T}, y::TracedRNumber{T}) where {T<:Real}
+function Base.complex(x::TracedRReal{T}, y::TracedRReal{T}) where {T<:Real}
     return @opcall complex(x, y)
 end
-Base.complex(x::TracedRNumber{T}) where {T<:Complex} = x
+Base.complex(x::TracedRComplex) = x
 
 for (jlop, hloop) in (
     (:(Base.min), :minimum),
@@ -247,20 +267,20 @@ for (jlop, hloop) in (
     end
 end
 
-function Base.rem(x::TracedRNumber, y::TracedRNumber, ::typeof(Base.RoundFromZero))
+function Base.rem(x::TracedRReal, y::TracedRReal, ::typeof(Base.RoundFromZero))
     return ifelse(
         signbit(x) == signbit(y), rem(x, y, Base.RoundUp), rem(x, y, Base.RoundDown)
     )
 end
 
-function Base.:*(x::TracedRNumber{T}, z::Complex{Bool}) where {T<:Real}
+function Base.:*(x::TracedRReal, z::Complex{Bool})
     # this is to support multiplication by im (Complex{Bool}(false, true))
     z_re, z_im = real(z), imag(z)
     res_re = z_re ? x : zero(x)
     res_im = z_im ? x : zero(x)
     return Complex(res_re, res_im)
 end
-Base.:*(z::Complex{Bool}, x::TracedRNumber{T}) where {T<:Real} = x * z
+Base.:*(z::Complex{Bool}, x::TracedRReal) = x * z
 
 # Based on https://github.com/JuliaLang/julia/blob/39255d47db7657950ff1c82137ecec5a70bae622/base/float.jl#L608-L617
 function Base.mod(
@@ -280,19 +300,19 @@ end
 for op in (:mod, :mod1, :rem)
     @eval begin
         function Base.$op(
-            @nospecialize(lhs::TracedRNumber{T}), @nospecialize(rhs::Number)
+            @nospecialize(lhs::TracedRReal{T}), @nospecialize(rhs::Real)
         ) where {T}
             return $(op)(lhs, Reactant.promote_to(TracedRNumber{T}, rhs))
         end
         function Base.$op(
-            @nospecialize(lhs::Number), @nospecialize(rhs::TracedRNumber{T})
+            @nospecialize(lhs::Real), @nospecialize(rhs::TracedRReal{T})
         ) where {T}
             return $(op)(Reactant.promote_to(TracedRNumber{T}, lhs), rhs)
         end
     end
 end
 
-Base.flipsign(x::TracedRNumber, y::TracedRNumber) = ifelse(y < 0, -x, x)
+Base.flipsign(x::TracedRReal, y::TracedRReal) = ifelse(y < 0, -x, x)
 
 function Base.div(
     x::TracedRNumber{<:Reactant.ReactantSInt}, y::TracedRNumber{<:Reactant.ReactantUInt}
@@ -346,20 +366,20 @@ function Base.div(
 end
 
 function Base.div(
-    @nospecialize(lhs::TracedRNumber{T}), rhs, r::Base.RoundingMode
-) where {T<:Integer}
+    @nospecialize(lhs::TracedRInteger{T}), rhs::Real, r::Base.RoundingMode
+) where {T}
     return div(lhs, Reactant.promote_to(TracedRNumber{T}, rhs), r)
 end
 function Base.div(
-    lhs, @nospecialize(rhs::TracedRNumber{T}), r::Base.RoundingMode
-) where {T<:Integer}
+    lhs::Real, @nospecialize(rhs::TracedRInteger{T}), r::Base.RoundingMode
+) where {T}
     return div(Reactant.promote_to(TracedRNumber{T}, lhs), rhs, r)
 end
 function Base.div(
-    @nospecialize(lhs::TracedRNumber{T1}),
-    @nospecialize(rhs::TracedRNumber{T2}),
+    @nospecialize(lhs::TracedRInteger{T1}),
+    @nospecialize(rhs::TracedRInteger{T2}),
     r::Base.RoundingMode,
-) where {T1<:Integer,T2<:Integer}
+) where {T1,T2}
     T = promote_type(T1, T2)
     return div(
         Reactant.promote_to(TracedRNumber{T}, lhs),
@@ -368,14 +388,14 @@ function Base.div(
     )
 end
 
-function Base.div(@nospecialize(lhs::TracedRNumber{T}), rhs) where {T}
+function Base.div(@nospecialize(lhs::TracedRReal{T}), rhs::Real) where {T}
     return @opcall divide(lhs, Reactant.promote_to(TracedRNumber{T}, rhs))
 end
-function Base.div(lhs, @nospecialize(rhs::TracedRNumber{T})) where {T}
+function Base.div(lhs::Real, @nospecialize(rhs::TracedRReal{T})) where {T}
     return @opcall divide(Reactant.promote_to(TracedRNumber{T}, lhs), rhs)
 end
 function Base.div(
-    @nospecialize(lhs::TracedRNumber{T1}), @nospecialize(rhs::TracedRNumber{T2})
+    @nospecialize(lhs::TracedRReal{T1}), @nospecialize(rhs::TracedRReal{T2})
 ) where {T1,T2}
     T = promote_type(T1, T2)
     return @opcall divide(
@@ -525,57 +545,77 @@ function Base.divrem(
 end
 
 function Base.:/(
-    @nospecialize(lhs::TracedRNumber{T}), @nospecialize(rhs::TracedRNumber{T})
-) where {T<:Integer}
+    @nospecialize(lhs::TracedRInteger{T}), @nospecialize(rhs::TracedRInteger{T})
+) where {T}
     return float(lhs) / float(rhs)
 end
 
-for (jlop, hloop, hlocomp) in (
-    (:(Base.:(==)), :compare, "EQ"),
-    (:(Base.:(!=)), :compare, "NE"),
-    (:(Base.:(>=)), :compare, "GE"),
-    (:(Base.:(>)), :compare, "GT"),
-    (:(Base.:(<=)), :compare, "LE"),
-    (:(Base.:(<)), :compare, "LT"),
-    (:(Base.isless), :compare, "LT"),
+# Comparisons are defined on the concrete traced types (not the `TracedRNumber`
+# union) so that they are strictly more specific than the corresponding Base
+# methods on `Real`/`AbstractFloat`/`Integer`. Order comparisons are real-only,
+# matching Base; `==`/`!=` also cover traced complex numbers. Mixed-type
+# arguments are promoted to a common traced type. The extra methods for
+# `Rational`, `BigInt`, `BigFloat`, `AbstractIrrational`, `AbstractFloat`, and
+# `Complex` resolve ambiguities with Base's specialized comparison methods.
+function promote_to_common(
+    @nospecialize(lhs::TracedRNumber{T1}), @nospecialize(rhs::TracedRNumber{T2})
+) where {T1,T2}
+    commonTy = TracedRNumber{Base.promote_type(T1, T2)}
+    return Reactant.promote_to(commonTy, lhs), Reactant.promote_to(commonTy, rhs)
+end
+function promote_to_common(@nospecialize(lhs::TracedRNumber), @nospecialize(rhs))
+    return lhs, Reactant.promote_to(lhs, rhs)
+end
+function promote_to_common(@nospecialize(lhs), @nospecialize(rhs::TracedRNumber))
+    return Reactant.promote_to(rhs, lhs), rhs
+end
+
+for (jlop, hlocomp, TracedTs) in (
+    (:(Base.:(==)), "EQ", (:TracedRInteger, :TracedRFloat, :TracedRComplex)),
+    (:(Base.:(!=)), "NE", (:TracedRInteger, :TracedRFloat, :TracedRComplex)),
+    (:(Base.:(>=)), "GE", (:TracedRInteger, :TracedRFloat)),
+    (:(Base.:(>)), "GT", (:TracedRInteger, :TracedRFloat)),
+    (:(Base.:(<=)), "LE", (:TracedRInteger, :TracedRFloat)),
+    (:(Base.:(<)), "LT", (:TracedRInteger, :TracedRFloat)),
+    (:(Base.isless), "LT", (:TracedRInteger, :TracedRFloat)),
 )
-    @eval begin
-        function $(jlop)(
-            @nospecialize(lhs::TracedRNumber{T}), @nospecialize(rhs::TracedRNumber{T})
-        ) where {T}
-            return @opcall compare(lhs, rhs; comparison_direction=$(hlocomp))
+    OtherT = length(TracedTs) == 2 ? :Real : :Number
+    disambTs = (:Rational, :BigInt, :BigFloat, :AbstractIrrational, :AbstractFloat)
+    length(TracedTs) == 3 && (disambTs = (disambTs..., :Complex))
+
+    for T in TracedTs
+        @eval begin
+            function $(jlop)(@nospecialize(lhs::$T{T}), @nospecialize(rhs::$T{T})) where {T}
+                return @opcall compare(lhs, rhs; comparison_direction=($(hlocomp)))
+            end
+            function $(jlop)(@nospecialize(lhs::$T), @nospecialize(rhs::$OtherT))
+                return $(jlop)(promote_to_common(lhs, rhs)...)
+            end
+            function $(jlop)(@nospecialize(lhs::$OtherT), @nospecialize(rhs::$T))
+                return $(jlop)(promote_to_common(lhs, rhs)...)
+            end
+        end
+        for X in disambTs
+            @eval begin
+                function $(jlop)(@nospecialize(lhs::$T), @nospecialize(rhs::$X))
+                    return $(jlop)(promote_to_common(lhs, rhs)...)
+                end
+                function $(jlop)(@nospecialize(lhs::$X), @nospecialize(rhs::$T))
+                    return $(jlop)(promote_to_common(lhs, rhs)...)
+                end
+            end
         end
+    end
+    for T1 in TracedTs, T2 in TracedTs
+        @eval function $(jlop)(@nospecialize(lhs::$T1), @nospecialize(rhs::$T2))
+            return $(jlop)(promote_to_common(lhs, rhs)...)
+        end
+    end
 
+    @eval begin
         # ambiguity fixes
         $(jlop)(@nospecialize(lhs::TracedRNumber), @nospecialize(::Missing)) = missing
         $(jlop)(@nospecialize(::Missing), @nospecialize(rhs::TracedRNumber)) = missing
-
-        function $(jlop)(@nospecialize(lhs::TracedRNumber{T}), @nospecialize(rhs)) where {T}
-            return $(jlop)(lhs, Reactant.promote_to(lhs, rhs))
-        end
-        function $(jlop)(
-            @nospecialize(lhs::TracedRNumber{T}), @nospecialize(rhs::Number)
-        ) where {T}
-            return $(jlop)(lhs, Reactant.promote_to(lhs, rhs))
-        end
-
-        function $(jlop)(@nospecialize(lhs), @nospecialize(rhs::TracedRNumber{T})) where {T}
-            return $(jlop)(Reactant.promote_to(rhs, lhs), rhs)
-        end
-        function $(jlop)(
-            @nospecialize(lhs::Number), @nospecialize(rhs::TracedRNumber{T})
-        ) where {T}
-            return $(jlop)(Reactant.promote_to(rhs, lhs), rhs)
-        end
-
-        function $(jlop)(
-            @nospecialize(lhs::TracedRNumber{T1}), @nospecialize(rhs::TracedRNumber{T2})
-        ) where {T1,T2}
-            commonTy = TracedRNumber{Base.promote_type(T1, T2)}
-            lhs = Reactant.promote_to(commonTy, lhs)
-            rhs = Reactant.promote_to(commonTy, rhs)
-            return $(jlop)(lhs, rhs)
-        end
     end
 end
 
@@ -618,7 +658,7 @@ function Base.ifelse(
 )
     @assert length(x) == length(y)
     ntuple(Val(length(x))) do i
-        Base.ifelse(pred, x[i], y[i])
+        return Base.ifelse(pred, x[i], y[i])
     end
 end
 
@@ -658,76 +698,32 @@ function Base.:*(x::TwicePrecision, v::TracedRNumber)
     return result
 end
 
-for (T1, T2) in Iterators.product((Bool, Integer), (Bool, Integer))
+for (jlop, hloop) in ((:(Base.:&), :and), (:(Base.:|), :or), (:(Base.xor), :xor))
     @eval begin
-        function Base.:&(x::TracedRNumber{A}, y::TracedRNumber{B}) where {A<:$(T1),B<:$(T2)}
-            C = promote_type(A, B)
-            return @opcall and(
-                Reactant.promote_to(TracedRNumber{C}, x),
-                Reactant.promote_to(TracedRNumber{C}, y),
-            )
-        end
-        function Base.:&(x::TracedRNumber{A}, y::B) where {A<:$(T1),B<:$(T2)}
-            C = promote_type(A, B)
-            return @opcall and(
-                Reactant.promote_to(TracedRNumber{C}, x),
-                Reactant.promote_to(TracedRNumber{C}, y),
-            )
-        end
-        function Base.:&(x::A, y::TracedRNumber{B}) where {A<:$(T1),B<:$(T2)}
-            C = promote_type(A, B)
-            return @opcall and(
-                Reactant.promote_to(TracedRNumber{C}, x),
-                Reactant.promote_to(TracedRNumber{C}, y),
-            )
-        end
-        function Base.:|(x::TracedRNumber{A}, y::TracedRNumber{B}) where {A<:$(T1),B<:$(T2)}
+        function $jlop(x::TracedRInteger{A}, y::TracedRInteger{B}) where {A,B}
             C = promote_type(A, B)
-            return @opcall or(
+            return @opcall $hloop(
                 Reactant.promote_to(TracedRNumber{C}, x),
                 Reactant.promote_to(TracedRNumber{C}, y),
             )
         end
-        function Base.:|(x::TracedRNumber{A}, y::B) where {A<:$(T1),B<:$(T2)}
-            C = promote_type(A, B)
-            return @opcall or(
+        function $jlop(x::TracedRInteger{A}, y::Integer) where {A}
+            C = promote_type(A, unwrapped_eltype(y))
+            return $jlop(
                 Reactant.promote_to(TracedRNumber{C}, x),
                 Reactant.promote_to(TracedRNumber{C}, y),
             )
         end
-        function Base.:|(x::A, y::TracedRNumber{B}) where {A<:$(T1),B<:$(T2)}
-            C = promote_type(A, B)
-            return @opcall or(
-                Reactant.promote_to(TracedRNumber{C}, x),
-                Reactant.promote_to(TracedRNumber{C}, y),
-            )
-        end
-        function Base.xor(
-            x::TracedRNumber{A}, y::TracedRNumber{B}
-        ) where {A<:$(T1),B<:$(T2)}
-            C = promote_type(A, B)
-            return @opcall xor(
-                Reactant.promote_to(TracedRNumber{C}, x),
-                Reactant.promote_to(TracedRNumber{C}, y),
-            )
-        end
-        function Base.xor(x::TracedRNumber{A}, y::B) where {A<:$(T1),B<:$(T2)}
-            C = promote_type(A, B)
-            return @opcall xor(
-                Reactant.promote_to(TracedRNumber{C}, x),
-                Reactant.promote_to(TracedRNumber{C}, y),
-            )
-        end
-        function Base.xor(x::A, y::TracedRNumber{B}) where {A<:$(T1),B<:$(T2)}
-            C = promote_type(A, B)
-            return @opcall xor(
+        function $jlop(x::Integer, y::TracedRInteger{B}) where {B}
+            C = promote_type(unwrapped_eltype(x), B)
+            return $jlop(
                 Reactant.promote_to(TracedRNumber{C}, x),
                 Reactant.promote_to(TracedRNumber{C}, y),
             )
         end
     end
 end
-Base.:!(x::TracedRNumber{<:Union{Bool,Integer}}) = @opcall not(x)
+Base.:!(x::TracedRInteger) = @opcall not(x)
 
 function Base.literal_pow(
     ::Base.RefValue{typeof(^)}, x::TracedRNumber{T}, ::Base.RefValue{Val{P}}
@@ -759,52 +755,57 @@ for (jlop, hloop) in (
     (:(Base.atan), :atan),
     (:(Base.atanh), :atanh),
     (:(Base.sign), :sign),
-    (:(Base.conj), :conj),
-    (:(Base.real), :real),
-    (:(Base.imag), :imag),
 )
-    @eval $(jlop)(@nospecialize(lhs::TracedRNumber)) = @opcall $(hloop)(lhs)
+    # Separate real and complex methods: the union would be ambiguous with
+    # Base's methods on `Real`.
+    @eval $(jlop)(@nospecialize(lhs::TracedRReal)) = @opcall $(hloop)(lhs)
+    @eval $(jlop)(@nospecialize(lhs::TracedRComplex)) = @opcall $(hloop)(lhs)
+end
+
+# `conj`, `real`, and `imag` of traced reals are covered by Base's identities.
+for (jlop, hloop) in ((:(Base.conj), :conj), (:(Base.real), :real), (:(Base.imag), :imag))
+    @eval $(jlop)(@nospecialize(lhs::TracedRComplex)) = @opcall $(hloop)(lhs)
 end
 
 # Degree-based trigonometric wrappers for TracedRNumber
 # These convert to radians internally so Reactant can lower to
 # StableHLO-supported radian trigonometric operations.
 
-Base.sind(x::TracedRNumber) = sin(deg2rad(x))
-Base.cosd(x::TracedRNumber) = cos(deg2rad(x))
-Base.tand(x::TracedRNumber) = tan(deg2rad(x))
-Base.cscd(x::TracedRNumber) = 1 / sind(x)
-Base.secd(x::TracedRNumber) = 1 / cosd(x)
-Base.cotd(x::TracedRNumber) = 1 / tand(x)
+Base.sind(x::TracedRReal) = sin(deg2rad(x))
+Base.cosd(x::TracedRReal) = cos(deg2rad(x))
+Base.tand(x::TracedRReal) = tan(deg2rad(x))
+Base.cscd(x::TracedRReal) = 1 / sind(x)
+Base.secd(x::TracedRReal) = 1 / cosd(x)
+Base.cotd(x::TracedRReal) = 1 / tand(x)
 
-Base.asind(x::TracedRNumber) = rad2deg(asin(x))
-Base.acosd(x::TracedRNumber) = rad2deg(acos(x))
-Base.atand(x::TracedRNumber) = rad2deg(atan(x))
+Base.asind(x::TracedRReal) = rad2deg(asin(x))
+Base.acosd(x::TracedRReal) = rad2deg(acos(x))
+Base.atand(x::TracedRReal) = rad2deg(atan(x))
 
-Base.atan(y::TracedRNumber, x::TracedRNumber) = @opcall atan2(y, x)
-Base.atand(y::TracedRNumber, x::TracedRNumber) = rad2deg(atan(y, x))
+Base.atan(y::TracedRReal{T}, x::TracedRReal{T}) where {T} = @opcall atan2(y, x)
+Base.atan(y::TracedRReal, x::TracedRReal) = atan(promote_to_common(y, x)...)
+Base.atand(y::TracedRReal, x::TracedRReal) = rad2deg(atan(y, x))
 
-Base.hypot(x::TracedRNumber{T}, y::TracedRNumber{T}) where {T} = @opcall hypot(x, y)
-Base.hypot(x::TracedRNumber{T}, y) where {T} = Base.hypot(x, Reactant.promote_to(x, y))
-function Base.hypot(x::TracedRNumber{T}, y::Number) where {T}
-    return Base.hypot(x, Reactant.promote_to(x, y))
+Base.hypot(x::TracedRReal{T}, y::TracedRReal{T}) where {T} = @opcall hypot(x, y)
+
+function Base.hypot(x::TracedRReal, y::Real)
+    return Base.hypot(promote_to_common(x, y)...)
 end
-Base.hypot(x, y::TracedRNumber{T}) where {T} = Base.hypot(Reactant.promote_to(y, x), y)
-function Base.hypot(x::Number, y::TracedRNumber{T}) where {T}
-    return Base.hypot(Reactant.promote_to(y, x), y)
+
+function Base.hypot(x::Real, y::TracedRReal)
+    return Base.hypot(promote_to_common(x, y)...)
 end
 
-function Base.hypot(x::TracedRNumber{T1}, y::TracedRNumber{T2}) where {T1,T2}
-    commonTy = TracedRNumber{Base.promote_type(T1, T2)}
-    return Base.hypot(Reactant.promote_to(commonTy, x), Reactant.promote_to(commonTy, y))
+function Base.hypot(x::TracedRReal, y::TracedRReal)
+    return Base.hypot(promote_to_common(x, y)...)
 end
 
 Base.hypot(::Missing, ::TracedRNumber) = missing
 Base.hypot(::TracedRNumber, ::Missing) = missing
 
-Base.acscd(x::TracedRNumber) = rad2deg(asin(1 / x))
-Base.asecd(x::TracedRNumber) = rad2deg(acos(1 / x))
-Base.acotd(x::TracedRNumber) = rad2deg(atan(1 / x))
+Base.acscd(x::TracedRReal) = rad2deg(asin(1 / x))
+Base.asecd(x::TracedRReal) = rad2deg(acos(1 / x))
+Base.acotd(x::TracedRReal) = rad2deg(atan(1 / x))
 
 for (jlop, hloop) in (
     (:(Base.sin), :sine),
@@ -827,26 +828,27 @@ for (jlop, hloop) in (
     (:(Base.atan), :atan),
     (:(Base.atanh), :atanh),
 )
-    @eval $(jlop)(@nospecialize(lhs::TracedRNumber{<:Integer})) =
-        @opcall $(hloop)(float(lhs))
+    @eval $(jlop)(@nospecialize(lhs::TracedRInteger)) = @opcall $(hloop)(float(lhs))
 end
 
 for (jlop, hloop) in
     ((:(Base.sinpi), :sine), (:(Base.cospi), :cosine), (:(Base.tanpi), :tan))
-    @eval $(jlop)(@nospecialize(lhs::TracedRNumber{T})) where {T} =
-        @opcall $(hloop)(T(π) * lhs)
+    for TracedT in (:TracedRReal, :TracedRComplex)
+        @eval $(jlop)(@nospecialize(lhs::$TracedT{T})) where {T} =
+            @opcall $(hloop)(T(π) * lhs)
+    end
 end
 
 function Base.sincospi(x::TracedRNumber{T}) where {T}
     return @opcall(sine(T(π) * x)), @opcall(cosine(T(π) * x))
 end
 
-function Base.cispi(x::TracedRNumber)
+function Base.cispi(x::TracedRReal)
     s, c = sincospi(x)
     return complex(c, s)
 end
 
-function Base.cis(x::TracedRNumber)
+function Base.cis(x::TracedRReal)
     s, c = sincos(x)
     return complex(c, s)
 end
@@ -856,7 +858,7 @@ function Base.sinc(x::TracedRNumber)
     return ifelse(isinf(x), zero(x), r)
 end
 
-function Base.sinc(x::TracedRNumber{<:Complex{T}}) where {T}
+function Base.sinc(x::TracedRComplex{Complex{T}}) where {T}
     r1 = ifelse(
         abs(x) < Base.Math._sinc_threshold(T),
         evalpoly(x^2, (T(1), -T(pi)^2 / 6, T(pi)^4 / 120)),
@@ -865,18 +867,17 @@ function Base.sinc(x::TracedRNumber{<:Complex{T}}) where {T}
     return ifelse(isinf(real(x)), zero(x), r1)
 end
 
-@noinline Base.Math.log10(x::TracedRNumber) = Base.Math._log(x, Val(10), :log10)
-@noinline Base.Math.log2(x::TracedRNumber) = Base.Math._log(x, Val(2), :log2)
+for TracedT in (:TracedRReal, :TracedRComplex)
+    @eval begin
+        @noinline Base.Math.log10(x::$TracedT) = Base.Math._log(x, Val(10), :log10)
+        @noinline Base.Math.log2(x::$TracedT) = Base.Math._log(x, Val(2), :log2)
+    end
+end
 Base.Math._log(x::TracedRNumber, base, ::Symbol) = log(x) / log(Reactant._unwrap_val(base))
 
-Base.isreal(x::TracedRNumber) = iszero(imag(x))
-Base.isreal(::TracedRNumber{<:Real}) = true
-
-Base.isinteger(x::TracedRNumber{<:Integer}) = true
-Base.isinteger(x::TracedRNumber{<:AbstractFloat}) = x - trunc(x) == zero(x)
-
-Base.isodd(x::TracedRNumber) = isodd(real(x))
-function Base.isodd(x::TracedRNumber{<:Real})
+Base.isodd(x::TracedRComplex) = isodd(real(x))
+Base.isodd(x::TracedRInteger) = !iszero(rem(x, 2))
+function Base.isodd(x::TracedRFloat)
     return (
         isinteger(x) &
         !iszero(
@@ -888,8 +889,9 @@ function Base.isodd(x::TracedRNumber{<:Real})
     )
 end
 
-Base.iseven(x::TracedRNumber) = iseven(real(x))
-function Base.iseven(x::TracedRNumber{<:Real})
+Base.iseven(x::TracedRComplex) = iseven(real(x))
+Base.iseven(x::TracedRInteger) = iszero(rem(x, 2))
+function Base.iseven(x::TracedRFloat)
     return (
         isinteger(x) & iszero(
             rem(
@@ -911,37 +913,41 @@ for (minT, maxT) in Iterators.product((Number, TracedRNumber), (Number, TracedRN
     end
 end
 
-function Base.float(x::TracedRNumber{T}) where {T}
+Base.float(x::TracedRFloat) = x
+function Base.float(x::Union{TracedRInteger{T},TracedRComplex{T}}) where {T}
     return Reactant.promote_to(TracedRNumber{float(T)}, x)
 end
 
+Base.float(::Type{TracedRNumber{T}}) where {T} = traced_number_type(float(T))
+for TracedT in (:TracedRInteger, :TracedRFloat, :TracedRComplex)
+    @eval Base.float(::Type{<:$TracedT{T}}) where {T} = traced_number_type(float(T))
+end
+
 using Reactant: ReactantFloat, ReactantInt
 
-Base.round(A::TracedRNumber{<:ReactantFloat}) = @opcall round_nearest_even(A)
-Base.round(A::TracedRNumber{<:ReactantInt}) = A
-function Base.round(A::TracedRNumber{<:ReactantFloat}, ::typeof(RoundNearest))
+Base.round(A::TracedRFloat) = @opcall round_nearest_even(A)
+Base.round(A::TracedRInteger) = A
+function Base.round(A::TracedRFloat, ::typeof(RoundNearest))
     return @opcall round_nearest_even(A)
 end
-function Base.round(A::TracedRNumber{<:ReactantFloat}, ::typeof(RoundNearestTiesAway))
+function Base.round(A::TracedRFloat, ::typeof(RoundNearestTiesAway))
     return @opcall round_nearest_afz(A)
 end
-Base.round(A::TracedRNumber{<:ReactantFloat}, ::typeof(RoundUp)) = ceil(A)
-Base.round(A::TracedRNumber{<:ReactantFloat}, ::typeof(RoundDown)) = floor(A)
-function Base.round(A::TracedRNumber{<:ReactantFloat}, ::typeof(RoundToZero))
+Base.round(A::TracedRFloat, ::typeof(RoundUp)) = ceil(A)
+Base.round(A::TracedRFloat, ::typeof(RoundDown)) = floor(A)
+function Base.round(A::TracedRFloat, ::typeof(RoundToZero))
     A_truncated = @opcall convert(TracedRNumber{Int}, A)
     return copysign(
         @opcall(convert(TracedRNumber{Reactant.unwrapped_eltype(A)}, A_truncated)), A
     )
 end
 
-Base.floor(A::TracedRNumber{<:ReactantFloat}) = @opcall floor(A)
-Base.floor(A::TracedRNumber{<:ReactantInt}) = A
-Base.ceil(A::TracedRNumber{<:ReactantFloat}) = @opcall ceil(A)
-Base.ceil(A::TracedRNumber{<:ReactantInt}) = A
+Base.floor(A::TracedRFloat) = @opcall floor(A)
+Base.floor(A::TracedRInteger) = A
+Base.ceil(A::TracedRFloat) = @opcall ceil(A)
+Base.ceil(A::TracedRInteger) = A
 
-function Base.unsafe_trunc(
-    T::Type{<:Reactant.ReactantInt}, x::TracedRNumber{<:Reactant.ReactantFloat}
-)
+function Base.unsafe_trunc(T::Type{<:Reactant.ReactantInt}, x::TracedRFloat)
     return @opcall convert(TracedRNumber{T}, x)
 end
 
@@ -952,7 +958,7 @@ for Ti in (Int8, Int16, Int32, Int64, Int128, UInt8, UInt16, UInt32, UInt64, UIn
             # directly. `Tf(typemax(Ti))+1` is either always exactly representable, or
             # rounded to `Inf` (e.g. when `Ti==UInt128 && Tf==Float32`).
             @eval begin
-                function Base.trunc(::Type{$Ti}, x::TracedRNumber{$Tf})
+                function Base.trunc(::Type{$Ti}, x::TracedRFloat{$Tf})
                     # TODO(#2236) throw error within traced
                     # if $(Tf(typemin(Ti))-one(Tf)) < x < $(Tf(typemax(Ti))+one(Tf))
                     return Base.unsafe_trunc($Ti, x)
@@ -967,7 +973,7 @@ for Ti in (Int8, Int16, Int32, Int64, Int128, UInt8, UInt16, UInt32, UInt64, UIn
             # be rounded up. This assumes that `Tf(typemin(Ti)) > -Inf`, which is true for
             # these types, but not for `Float16` or larger integer types.
             @eval begin
-                function Base.trunc(::Type{$Ti}, x::TracedRNumber{$Tf})
+                function Base.trunc(::Type{$Ti}, x::TracedRFloat{$Tf})
                     # TODO(#2236) throw error within traced
                     # if $(Tf(typemin(Ti))) <= x < $(Tf(typemax(Ti)))
                     return Base.unsafe_trunc($Ti, x)
@@ -982,9 +988,9 @@ end
 
 # matches convert methods
 # also determines floor, ceil, round
-Base.trunc(::Type{Signed}, x::TracedRNumber{<:Base.IEEEFloat}) = Base.trunc(Int, x)
-Base.trunc(::Type{Unsigned}, x::TracedRNumber{<:Base.IEEEFloat}) = Base.trunc(UInt, x)
-Base.trunc(::Type{Integer}, x::TracedRNumber{<:Base.IEEEFloat}) = Base.trunc(Int, x)
+Base.trunc(::Type{Signed}, x::TracedRFloat{<:Base.IEEEFloat}) = Base.trunc(Int, x)
+Base.trunc(::Type{Unsigned}, x::TracedRFloat{<:Base.IEEEFloat}) = Base.trunc(UInt, x)
+Base.trunc(::Type{Integer}, x::TracedRFloat{<:Base.IEEEFloat}) = Base.trunc(Int, x)
 
 function (::Type{T})(x::TwicePrecision) where {T<:Reactant.TracedRNumber}
     return (T(x.hi) + T(x.lo))::T
@@ -994,13 +1000,13 @@ function (::Type{T})(x::TwicePrecision) where {T<:Reactant.ConcreteRNumber}
     return Reactant.ConcreteRNumber(T(x.hi) - T(x.lo))::T
 end
 
-function Base.round(::Type{T}, x::TracedRNumber{<:AbstractFloat}) where {T<:Integer}
+function Base.round(::Type{T}, x::TracedRFloat) where {T<:Integer}
     return trunc(T, round(x))
 end
-function Base.floor(::Type{T}, x::TracedRNumber{<:AbstractFloat}) where {T<:Integer}
+function Base.floor(::Type{T}, x::TracedRFloat) where {T<:Integer}
     return trunc(T, floor(x))
 end
-function Base.ceil(::Type{T}, x::TracedRNumber{<:AbstractFloat}) where {T<:Integer}
+function Base.ceil(::Type{T}, x::TracedRFloat) where {T<:Integer}
     return trunc(T, ceil(x))
 end
 
@@ -1035,22 +1041,22 @@ function Base.typed_hvncat(
     return Base.typed_hvncat(T, dims, row_first, xs...)
 end
 
+Base.signbit(x::TracedRInteger{<:Reactant.ReactantSInt}) = x < 0
+function Base.signbit(::TracedRInteger{<:Union{Bool,Reactant.ReactantUInt}})
+    return Reactant.promote_to(TracedRNumber{Bool}, false)
+end
 for (Ti, Tf) in ((Int16, Float16), (Int32, Float32), (Int64, Float64))
-    @eval begin
-        Base.signbit(x::TracedRNumber{$(Ti)}) = x < 0
-        Base.signbit(x::TracedRNumber{$(Tf)}) = signbit(@opcall(bitcast_convert($(Ti), x)))
-    end
+    @eval Base.signbit(x::TracedRFloat{$(Tf)}) = signbit(@opcall(bitcast_convert($(Ti), x)))
 end
-Base.signbit(::TracedRNumber{<:Unsigned}) = Reactant.promote_to(TracedRNumber{Bool}, false)
 
-function Base.copysign(x::TracedRNumber, y::TracedRNumber)
+function Base.copysign(x::TracedRReal, y::TracedRReal)
     return ifelse(signbit(y), -one(x), one(x)) * abs(x)
 end
-function Base.copysign(x::TracedRNumber{T}, y::S) where {T,S<:Number}
-    return copysign(x, Reactant.promote_to(TracedRNumber{S}, y))
+function Base.copysign(x::TracedRReal, y::S) where {S<:Real}
+    return copysign(x, Reactant.promote_to(TracedRNumber{unwrapped_eltype(S)}, y))
 end
-function Base.copysign(x::S, y::TracedRNumber{T}) where {S<:Number,T}
-    return copysign(Reactant.promote_to(TracedRNumber{S}, x), y)
+function Base.copysign(x::S, y::TracedRReal) where {S<:Real}
+    return copysign(Reactant.promote_to(TracedRNumber{unwrapped_eltype(S)}, x), y)
 end
 
 function Base.zeros(::Type{<:TracedRNumber{T}}, dims::Dims{N}) where {T,N}
@@ -1090,44 +1096,42 @@ end
 # rem2pi: fallback to rem for now.
 # TODO(#2259): we should replace with the more numerically stable version.
 # https://github.com/JuliaLang/julia/blob/4d04bb6b3b1b879f4dbb918d194c5c939a1e7f3c/base/special/rem2pi.jl#L133
-Base.rem2pi(x::TracedRNumber, r::Base.RoundingMode) = rem(x, typeof(x)(2π), r)
+Base.rem2pi(x::TracedRReal, r::Base.RoundingMode) = rem(x, typeof(x)(2π), r)
 
 function Base.rem2pi(
     x::T, r::Base.RoundingMode
-) where {T<:TracedRNumber{<:Union{Float16,Float32}}}
+) where {T<:TracedRFloat{<:Union{Float16,Float32}}}
     return T(rem2pi(TracedRNumber{Float64}(x), r))
 end
-Base.rem2pi(x::TracedRNumber{<:Integer}, r::Base.RoundingMode) = rem2pi(float(x), r)
+Base.rem2pi(x::TracedRInteger, r::Base.RoundingMode) = rem2pi(float(x), r)
 
 @static if isdefined(Base, :unchecked_oneto)
-    function Base.unchecked_oneto(x::TracedRNumber{<:Integer})
+    function Base.unchecked_oneto(x::TracedRInteger)
         return Reactant.TracedUnitRange(one(x), x)
     end
 end
 
-Base.unsigned(x::TracedRNumber{T}) where {T<:Reactant.ReactantUInt} = x
-function Base.unsigned(x::TracedRNumber{T}) where {T<:Reactant.ReactantSInt}
+Base.unsigned(x::TracedRInteger{T}) where {T<:Reactant.ReactantUInt} = x
+function Base.unsigned(x::TracedRInteger{T}) where {T<:Reactant.ReactantSInt}
     return convert(TracedRNumber{unsigned(T)}, x)
 end
 
-function Base.signed(x::TracedRNumber{T}) where {T<:Reactant.ReactantUInt}
+function Base.signed(x::TracedRInteger{T}) where {T<:Reactant.ReactantUInt}
     return convert(TracedRNumber{signed(T)}, x)
 end
-Base.signed(x::TracedRNumber{T}) where {T<:Reactant.ReactantSInt} = x
+Base.signed(x::TracedRInteger{T}) where {T<:Reactant.ReactantSInt} = x
 
-Base.uabs(x::TracedRNumber{<:Integer}) = abs(x)
-Base.uabs(x::TracedRNumber{<:Reactant.ReactantSInt}) = unsigned(abs(x))
+Base.uabs(x::TracedRInteger) = abs(x)
+Base.uabs(x::TracedRInteger{<:Reactant.ReactantSInt}) = unsigned(abs(x))
 
-function Base.divgcd(
-    x::TracedRNumber{TX}, y::TracedRNumber{TY}
-) where {TX<:Integer,TY<:Integer}
+function Base.divgcd(x::TracedRInteger, y::TracedRInteger)
     g = gcd(Base.uabs(x), Base.uabs(y))
     return div(x, g), div(y, g)
 end
 
 # See https://github.com/EnzymeAD/Reactant.jl/issues/2476 for why this is
 # implemented this way.
-function Base.gcd(a::TracedRNumber{T}, b::TracedRNumber{T}) where {T<:Integer}
+function Base.gcd(a::TracedRInteger{T}, b::TracedRInteger{T}) where {T}
     a_tmp = copy(a)
     b_tmp = copy(b)
     @trace while b_tmp != 0

From a31035e68e11814f88ae1898e4a437d688590578 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Tue, 7 Jul 2026 20:13:16 +0200
Subject: [PATCH 05/30] Adapt type tracing to the split number types

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.
---
 src/Tracing.jl | 114 ++++++++++++++++++++++++++++++++-----------------
 1 file changed, 74 insertions(+), 40 deletions(-)

diff --git a/src/Tracing.jl b/src/Tracing.jl
index 288020cfa9..7fc502826c 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -37,6 +37,11 @@ for T in (
     AbstractFloat,
     Integer,
     RNumber,
+    # The abstract kinds resolve the ambiguity between the `AbstractFloat` /
+    # `Integer` methods above and the `RNumber` union.
+    RInteger,
+    RFloat,
+    RComplex,
     Val,
     VersionNumber,
     Sharding.Mesh,
@@ -63,15 +68,15 @@ Base.@nospecializeinfer function traced_type_inner(
 )
     if mode == ArrayToConcrete && T <: track_numbers
         if runtime isa Val{:PJRT}
-            return ConcretePJRTNumber{T,_unwrap_val(ndevices)}
+            return pjrt_number_type(T){_unwrap_val(ndevices)}
         elseif runtime isa Val{:IFRT}
-            return ConcreteIFRTNumber{T}
+            return ifrt_number_type(T)
         else
             error("Unsupported runtime $runtime")
         end
     elseif (mode == NoStopTracedTrack || mode == TracedTrack || mode == TracedSetPath) &&
         T <: track_numbers
-        return TracedRNumber{T}
+        return traced_number_type(T)
     end
     return T
 end
@@ -221,54 +226,71 @@ Base.@nospecializeinfer function traced_type_inner(
     end
 end
 
-Base.@nospecializeinfer function traced_type_inner(
-    @nospecialize(T0::Type{<:ConcretePJRTNumber}),
-    seen,
-    @nospecialize(mode::TraceMode),
-    @nospecialize(track_numbers::Type),
-    @nospecialize(ndevices),
-    @nospecialize(runtime)
+Base.@nospecializeinfer function concrete_pjrt_number_traced_type(
+    @nospecialize(T0::Type), @nospecialize(mode::TraceMode), @nospecialize(ndevices)
 )
-    T = Base.unwrap_unionall(T0).parameters[1]
-
     if mode == ConcreteToTraced
-        return T0 isa UnionAll ? TracedRNumber : TracedRNumber{T}
+        T0 isa UnionAll && return TracedRNumber
+        return traced_number_type(unwrapped_eltype(T0))
     elseif mode == TracedToConcrete
         return T0
     elseif mode == ArrayToConcrete
-        @assert runtime isa Val{:PJRT}
         if T0 isa UnionAll
-            return ConcretePJRTNumbe{T,_unwrap_val(ndevices)} where {T}
+            return ConcretePJRTNumber{T,_unwrap_val(ndevices)} where {T}
         else
-            return ConcretePJRTNumber{T,_unwrap_val(ndevices)}
+            return pjrt_number_type(unwrapped_eltype(T0)){_unwrap_val(ndevices)}
         end
     else
         throw("Unsupported mode: $mode")
     end
 end
 
-Base.@nospecializeinfer function traced_type_inner(
-    @nospecialize(T0::Type{<:ConcreteIFRTNumber}),
-    seen,
-    @nospecialize(mode::TraceMode),
-    @nospecialize(track_numbers::Type),
-    @nospecialize(ndevices),
-    @nospecialize(runtime)
+Base.@nospecializeinfer function concrete_ifrt_number_traced_type(
+    @nospecialize(T0::Type), @nospecialize(mode::TraceMode)
 )
-    T = Base.unwrap_unionall(T0).parameters[1]
-
     if mode == ConcreteToTraced
-        return T0 isa UnionAll ? TracedRNumber : TracedRNumber{T}
+        T0 isa UnionAll && return TracedRNumber
+        return traced_number_type(unwrapped_eltype(T0))
     elseif mode == TracedToConcrete
         return T0
     elseif mode == ArrayToConcrete
-        @assert runtime isa Val{:IFRT}
-        return T0 isa UnionAll ? ConcreteIFRTNumber : ConcreteIFRTNumber{T}
+        T0 isa UnionAll && return ConcreteIFRTNumber
+        return ifrt_number_type(unwrapped_eltype(T0))
     else
         throw("Unsupported mode: $mode")
     end
 end
 
+for NumT in
+    (:ConcretePJRTNumber, :ConcretePJRTInteger, :ConcretePJRTFloat, :ConcretePJRTComplex)
+    @eval Base.@nospecializeinfer function traced_type_inner(
+        @nospecialize(T0::Type{<:$NumT}),
+        seen,
+        @nospecialize(mode::TraceMode),
+        @nospecialize(track_numbers::Type),
+        @nospecialize(ndevices),
+        @nospecialize(runtime)
+    )
+        mode == ArrayToConcrete && @assert runtime isa Val{:PJRT}
+        return concrete_pjrt_number_traced_type(T0, mode, ndevices)
+    end
+end
+
+for NumT in
+    (:ConcreteIFRTNumber, :ConcreteIFRTInteger, :ConcreteIFRTFloat, :ConcreteIFRTComplex)
+    @eval Base.@nospecializeinfer function traced_type_inner(
+        @nospecialize(T0::Type{<:$NumT}),
+        seen,
+        @nospecialize(mode::TraceMode),
+        @nospecialize(track_numbers::Type),
+        @nospecialize(ndevices),
+        @nospecialize(runtime)
+    )
+        mode == ArrayToConcrete && @assert runtime isa Val{:IFRT}
+        return concrete_ifrt_number_traced_type(T0, mode)
+    end
+end
+
 Base.@nospecializeinfer function traced_type_inner(
     @nospecialize(A::Type{<:ConcretePJRTArray}),
     seen,
@@ -351,12 +373,13 @@ Base.@nospecializeinfer function traced_type_inner(
     if mode == ConcreteToTraced
         throw("TracedRArray cannot be traced")
     elseif mode == TracedToConcrete
+        elT, N = Base.unwrap_unionall(T).parameters
         if runtime isa Val{:PJRT}
-            return ConcretePJRTArray{T.parameters[1],T.parameters[2],_unwrap_val(ndevices)}
+            return ConcretePJRTArray{elT,N,_unwrap_val(ndevices)}
         elseif runtime isa Val{:IFRT}
             return ConcreteIFRTArray{
-                T.parameters[1],
-                T.parameters[2],
+                elT,
+                N,
                 Nothing, # TODO(#2251): check if we can ensure no padding??
             }
         end
@@ -370,11 +393,9 @@ Base.@nospecializeinfer function traced_type_inner(
     end
 end
 
-Base.@nospecializeinfer function traced_type_inner(
-    @nospecialize(T::Type{<:TracedRNumber}),
-    seen,
-    mode::TraceMode,
-    @nospecialize(track_numbers::Type),
+Base.@nospecializeinfer function traced_number_traced_type(
+    @nospecialize(T::Type),
+    @nospecialize(mode::TraceMode),
     @nospecialize(ndevices),
     @nospecialize(runtime)
 )
@@ -385,12 +406,12 @@ Base.@nospecializeinfer function traced_type_inner(
             if T isa UnionAll
                 return UnionAll(T.var, ConcretePJRTNumber{T.var,_unwrap_val(ndevices)})
             end
-            return ConcretePJRTNumber{T.parameters[1],_unwrap_val(ndevices)}
+            return pjrt_number_type(unwrapped_eltype(T)){_unwrap_val(ndevices)}
         elseif runtime isa Val{:IFRT}
             if T isa UnionAll
                 return UnionAll(T.var, ConcreteIFRTNumber{T.var})
             end
-            return ConcreteIFRTNumber{T.parameters[1]}
+            return ifrt_number_type(unwrapped_eltype(T))
         end
         error("Unsupported runtime $runtime")
     elseif mode == TracedToJAX
@@ -402,6 +423,19 @@ Base.@nospecializeinfer function traced_type_inner(
     end
 end
 
+for NumT in (:TracedRNumber, :TracedRInteger, :TracedRFloat, :TracedRComplex)
+    @eval Base.@nospecializeinfer function traced_type_inner(
+        @nospecialize(T::Type{<:$NumT}),
+        seen,
+        mode::TraceMode,
+        @nospecialize(track_numbers::Type),
+        @nospecialize(ndevices),
+        @nospecialize(runtime)
+    )
+        return traced_number_traced_type(T, mode, ndevices, runtime)
+    end
+end
+
 Base.@nospecializeinfer function traced_type_inner(
     @nospecialize(A::Type{AbstractArray}),
     seen,
@@ -950,7 +984,7 @@ Reactant.apply_type_with_promotion(Foo, (Int, TracedRArray{Int, 1}))
 ```
 returns
 ```jl
-(Foo{Reactant.TracedRNumber{Int64}, Reactant.TracedRArray{Int64, 1}}, Bool[1, 0])
+(Foo{Reactant.TracedRInteger{Int64}, Reactant.TracedRArray{Int64, 1}}, Bool[1, 0])
 ```
 
 The first type parameter has been promoted to satisfy to be in agreement with the second parameter.
@@ -1002,7 +1036,7 @@ function apply_type_with_promotion(wrapper, params, relevant_typevars=typevar_di
                         # This happens when `value` lost the promotion battle.
                         # At this point, we need to update the problematic parameter in`value`.
                         d = typevar_dict(rewrapped)
-                        v = Any[param.parameters...]
+                        v = Any[Base.unwrap_unionall(param).parameters...]
                         v[d[typevar]] = resolved
                         params[i], _changed_params = apply_type_with_promotion(rewrapped, v)
                     end

From 337242ea0c704e33e20e3190036797f15f422f91 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Tue, 7 Jul 2026 20:13:16 +0200
Subject: [PATCH 06/30] Adapt traced arrays, ops, ranges, and rationals to the
 split number types

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.
---
 src/Indexing.jl       |   6 +-
 src/Ops.jl            | 189 +++++++++++++++++++++---------------------
 src/TracedRArray.jl   |   4 +-
 src/TracedRange.jl    |  15 ++--
 src/TracedRational.jl |  33 ++++----
 src/TracedUtils.jl    |  29 +++----
 6 files changed, 135 insertions(+), 141 deletions(-)

diff --git a/src/Indexing.jl b/src/Indexing.jl
index e3f664baf9..3b2156216b 100644
--- a/src/Indexing.jl
+++ b/src/Indexing.jl
@@ -203,7 +203,7 @@ function Base.getindex(
 end
 
 function Base.getindex(
-    x::Base.Sort.WithoutMissingVector{TracedRNumber{T}}, i::Int
+    x::Base.Sort.WithoutMissingVector{<:TracedRNumber{T}}, i::Int
 ) where {T}
     out = getindex(x.data, i)
     @assert !(out isa Missing)
@@ -442,9 +442,7 @@ function Base.setindex!(a::TracedRArray{T,N}, v, indices::Vararg{Any,N}) where {
 
     TracedUtils.set_mlir_data!(
         a,
-        @opcall(
-            dynamic_update_slice(a, v, [i isa Colon ? 1 : first(i) for i in indices])
-        ).mlir_data,
+        @opcall(dynamic_update_slice(a, v, [i isa Colon ? 1 : first(i) for i in indices])).mlir_data,
     )
     return v
 end
diff --git a/src/Ops.jl b/src/Ops.jl
index 883cce7b46..807df960ff 100644
--- a/src/Ops.jl
+++ b/src/Ops.jl
@@ -41,16 +41,14 @@ const SVD_ALGORITHM_MAP = Dict(
 )
 
 function _function_macro_error()
-    throw(ArgumentError("`caller_function` is not available in this context"))
+    return throw(ArgumentError("`caller_function` is not available in this context"))
 end
 
 macro caller_function()
-    return esc(
-        quote
-            $(Expr(:isdefined, :var"#self#")) || $(_function_macro_error)()
-            var"#self#"
-        end,
-    )
+    return esc(quote
+        $(Expr(:isdefined, :var"#self#")) || $(_function_macro_error)()
+        var"#self#"
+    end)
 end
 
 """
@@ -121,6 +119,11 @@ function mlir_type(RT::Type{<:RArray{T,N}}, shape) where {T,N}
     return MLIR.IR.TensorType(collect(Int, shape), MLIR.IR.Type(unwrapped_eltype(RT)))
 end
 
+function mlir_type(RT::Type{<:TracedRArray{T,N}}, shape) where {T,N}
+    @assert length(shape) == N
+    return MLIR.IR.TensorType(collect(Int, shape), MLIR.IR.Type(T))
+end
+
 function mlir_type(RT::Type{<:RNumber})::MLIR.IR.Type
     return MLIR.IR.TensorType(Int[], MLIR.IR.Type(unwrapped_eltype(RT)))
 end
@@ -1374,16 +1377,15 @@ end
         sample_inputs[2i - 1] = Reactant.promote_to(TracedRNumber{T}, 0)
         sample_inputs[2i] = Reactant.promote_to(TracedRNumber{T}, 0)
     end
-    func =
-        Reactant.TracedUtils.make_mlir_fn(
-            comparator,
-            (sample_inputs...,),
-            (),
-            "comparator",
-            false;
-            args_in_result=:none,
-            return_dialect=:stablehlo,
-        ).f
+    func = Reactant.TracedUtils.make_mlir_fn(
+        comparator,
+        (sample_inputs...,),
+        (),
+        "comparator",
+        false;
+        args_in_result=:none,
+        return_dialect=:stablehlo,
+    ).f
     @assert MLIR.IR.nregions(func) == 1
     fn_name = String(
         MLIR.IR.getattr(func, String(MLIR.API.mlirSymbolTableGetSymbolAttributeName()))
@@ -1432,21 +1434,20 @@ end
 ) where {T<:AbstractFloat,N}
     fallback === missing && (fallback = Reactant.FALLBACK_APPROX_TOP_K_LOWERING[])
 
-    func =
-        Reactant.TracedUtils.make_mlir_fn(
-            comparator,
-            (
-                Reactant.promote_to(TracedRNumber{T}, 0),
-                Reactant.promote_to(TracedRNumber{T}, 0),
-                Reactant.promote_to(TracedRNumber{Int32}, 0),
-                Reactant.promote_to(TracedRNumber{Int32}, 0),
-            ),
-            (),
-            "comparator",
-            false;
-            args_in_result=:none,
-            return_dialect=:stablehlo,
-        ).f
+    func = Reactant.TracedUtils.make_mlir_fn(
+        comparator,
+        (
+            Reactant.promote_to(TracedRNumber{T}, 0),
+            Reactant.promote_to(TracedRNumber{T}, 0),
+            Reactant.promote_to(TracedRNumber{Int32}, 0),
+            Reactant.promote_to(TracedRNumber{Int32}, 0),
+        ),
+        (),
+        "comparator",
+        false;
+        args_in_result=:none,
+        return_dialect=:stablehlo,
+    ).f
     @assert MLIR.IR.nregions(func) == 1
     fn_name = MLIR.IR.FlatSymbolRefAttribute(
         String(
@@ -1504,7 +1505,7 @@ end
     @assert 1 <= dimension <= N
 
     # XLA codegen for top.k is extremely sub-optimal. For special cases we can bypass that
-    if k isa Integer && k == 1
+    if k isa Integer && !(k isa TracedRNumber) && k == 1
         values, indices = argmax(x; dimension, location)
         return (;
             values, indices=add(indices, fill(Int64(1), Tuple(size(indices))); location)
@@ -2177,16 +2178,15 @@ end
         Reactant.promote_to(TracedRNumber, zero(T)),
     )
 
-    compiled_fn =
-        Reactant.TracedUtils.make_mlir_fn(
-            f,
-            sample_inputs,
-            (),
-            "update_computation",
-            false;
-            args_in_result=:result,
-            return_dialect=:stablehlo,
-        ).f
+    compiled_fn = Reactant.TracedUtils.make_mlir_fn(
+        f,
+        sample_inputs,
+        (),
+        "update_computation",
+        false;
+        args_in_result=:result,
+        return_dialect=:stablehlo,
+    ).f
     update_computation = MLIR.IR.Region()
     MLIR.API.mlirRegionTakeBody(update_computation, MLIR.IR.region(compiled_fn, 1))
     MLIR.IR.rmfromparent!(compiled_fn)
@@ -2370,36 +2370,34 @@ end
 
     input_types = [mlir_type(arg) for arg in linear_args]
 
-    cond_fn_compiled =
-        Reactant.TracedUtils.make_mlir_fn(
-            cond_fn,
-            traced_args,
-            (),
-            string(gensym("cond_fn")),
-            false;
-            return_dialect=:stablehlo,
-            args_in_result=:result,
-            do_transpose=false,
-            argprefix=gensym("loop_condarg"),
-            resprefix=gensym("loop_condres"),
-            resargprefix=gensym("loop_condresarg"),
-        ).f
-
-    body_fn_compiled =
-        Reactant.TracedUtils.make_mlir_fn(
-            body_fn,
-            traced_args,
-            (),
-            string(gensym("body_fn")),
-            false;
-            return_dialect=:stablehlo,
-            args_in_result=:all,
-            do_transpose=false,
-            verify_arg_names,
-            argprefix=gensym("loop_bodyarg"),
-            resprefix=gensym("loop_bodyres"),
-            resargprefix=gensym("loop_bodyresarg"),
-        ).f
+    cond_fn_compiled = Reactant.TracedUtils.make_mlir_fn(
+        cond_fn,
+        traced_args,
+        (),
+        string(gensym("cond_fn")),
+        false;
+        return_dialect=:stablehlo,
+        args_in_result=:result,
+        do_transpose=false,
+        argprefix=gensym("loop_condarg"),
+        resprefix=gensym("loop_condres"),
+        resargprefix=gensym("loop_condresarg"),
+    ).f
+
+    body_fn_compiled = Reactant.TracedUtils.make_mlir_fn(
+        body_fn,
+        traced_args,
+        (),
+        string(gensym("body_fn")),
+        false;
+        return_dialect=:stablehlo,
+        args_in_result=:all,
+        do_transpose=false,
+        verify_arg_names,
+        argprefix=gensym("loop_bodyarg"),
+        resprefix=gensym("loop_bodyres"),
+        resargprefix=gensym("loop_bodyresarg"),
+    ).f
 
     cond_reg = Reactant.TracedUtils.__take_region(cond_fn_compiled)
     body_reg = Reactant.TracedUtils.__take_region(body_fn_compiled)
@@ -2844,16 +2842,17 @@ end
         location,
     )
 
-    corrected_traced_results =
-        map(zip(traced_false_results, traced_true_results)) do (fr, tr)
-            if fr isa MissingTracedValue && tr isa MissingTracedValue
-                return fr
-            elseif fr isa MissingTracedValue
-                return tr
-            else
-                return fr
-            end
+    corrected_traced_results = map(
+        zip(traced_false_results, traced_true_results)
+    ) do (fr, tr)
+        if fr isa MissingTracedValue && tr isa MissingTracedValue
+            return fr
+        elseif fr isa MissingTracedValue
+            return tr
+        else
+            return fr
         end
+    end
 
     @assert length(all_paths) == length(result_types)
 
@@ -3492,16 +3491,15 @@ end
 function _construct_reduce_function(f::F, Ts::Type...) where {F}
     inputs_1 = [Reactant.promote_to(TracedRNumber{T}, 0) for T in Ts]
     inputs_2 = [Reactant.promote_to(TracedRNumber{T}, 0) for T in Ts]
-    func =
-        Reactant.TracedUtils.make_mlir_fn(
-            f,
-            (inputs_1..., inputs_2...),
-            (),
-            "reduce_fn" * string(f),
-            false;
-            args_in_result=:none,
-            return_dialect=:stablehlo,
-        ).f
+    func = Reactant.TracedUtils.make_mlir_fn(
+        f,
+        (inputs_1..., inputs_2...),
+        (),
+        "reduce_fn" * string(f),
+        false;
+        args_in_result=:none,
+        return_dialect=:stablehlo,
+    ).f
 
     @assert MLIR.IR.nregions(func) == 1
     ftype_attr = MLIR.IR.getattr(func, "function_type")
@@ -3628,14 +3626,15 @@ function standardize_start_index(
     start_index::Union{Integer,TracedRNumber{<:Integer}},
     idx::Integer,
 )
-    if (start_index isa Integer && start_index ≤ typemax(Int32)) || sz ≤ typemax(Int32)
-        if start_index isa Integer && update_sz !== nothing
+    is_static_index = !(start_index isa TracedRNumber)
+    if (is_static_index && start_index ≤ typemax(Int32)) || sz ≤ typemax(Int32)
+        if is_static_index && update_sz !== nothing
             @assert start_index + update_sz - 1 ≤ sz "Index $(idx) out of bounds: \
                                                   start_index=$(start_index), \
                                                   update_sz=$(update_sz), sz=$(sz)"
         end
         start_index = Reactant.promote_to(TracedRNumber{Int32}, start_index)
-    elseif start_index isa Integer && update_sz !== nothing
+    elseif is_static_index && update_sz !== nothing
         @assert start_index + update_sz - 1 ≤ sz "Index $(idx) out of bounds: \
                                               start_index=$(start_index), \
                                               update_sz=$(update_sz), sz=$(sz)"
diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
index 0d4144753c..77a36e2d70 100644
--- a/src/TracedRArray.jl
+++ b/src/TracedRArray.jl
@@ -252,7 +252,7 @@ function Base.fill!(A::AnyTracedRArray{T,N}, x::TracedRNumber{T2}) where {T,N,T2
 end
 
 function Base.fill!(A::Array{T,N}, x::TracedRNumber{T2}) where {T,N,T2}
-    throw(MethodError(fill!, (A, x)))
+    return throw(MethodError(fill!, (A, x)))
 end
 
 struct AbstractReactantArrayStyle{N} <: AbstractArrayStyle{N} end
@@ -1007,7 +1007,7 @@ end
         end
         return TracedRNumber{
             unwrapped_eltype(
-                Base._accumulate_promote_op(op, Array{T,ndims(A)}(undef, size(A)); init)
+                Base._accumulate_promote_op(op, Array{T,ndims(A)}(undef, size(A));init)
             ),
         }
     end
diff --git a/src/TracedRange.jl b/src/TracedRange.jl
index 1e8fa1ed93..3b74fd559d 100644
--- a/src/TracedRange.jl
+++ b/src/TracedRange.jl
@@ -1,6 +1,6 @@
 module TracedRangeOverrides
 
-using ..Reactant: Reactant, AbstractConcreteNumber, TracedRNumber
+using ..Reactant: Reactant, AbstractConcreteNumber, TracedRNumber, TracedRReal
 import ..Reactant: TracedStepRangeLen, TracedUnitRange
 using ..Ops: @opcall
 
@@ -182,15 +182,14 @@ function Base._reshape(parent::TracedUnitRange, dims::Dims)
     return Base.__reshape((parent, IndexStyle(parent)), dims)
 end
 
-function (C::Base.Colon)(start::TracedRNumber{T}, stop::TracedRNumber{T}) where {T}
+# These need to dominate Base's promoting `(:)(::Real, ::Real)` so that ranges
+# with traced endpoints become `TracedUnitRange` instead of Base ranges.
+function (C::Base.Colon)(start::TracedRReal{T}, stop::TracedRReal{T}) where {T}
     return TracedUnitRange(start, stop)
 end
-function (C::Base.Colon)(start::TracedRNumber{T}, stop::T) where {T}
-    return C(start, TracedRNumber{T}(stop))
-end
-function (C::Base.Colon)(start::T, stop::TracedRNumber{T}) where {T}
-    return C(TracedRNumber{T}(start), stop)
-end
+(C::Base.Colon)(start::TracedRReal, stop::TracedRReal) = C(promote(start, stop)...)
+(C::Base.Colon)(start::TracedRReal, stop::Real) = C(promote(start, stop)...)
+(C::Base.Colon)(start::Real, stop::TracedRReal) = C(promote(start, stop)...)
 
 Base.maximum(r::TracedUnitRange) = last(r)
 
diff --git a/src/TracedRational.jl b/src/TracedRational.jl
index f8baa1a794..88b80fe2fb 100644
--- a/src/TracedRational.jl
+++ b/src/TracedRational.jl
@@ -1,6 +1,6 @@
 module TracedRationalOverrides
 
-using ..Reactant: Reactant, AbstractConcreteNumber, TracedRNumber
+using ..Reactant: Reactant, AbstractConcreteNumber, TracedRNumber, TracedRInteger, RInteger
 import ..Reactant: TracedRational
 using ReactantCore: ReactantCore
 
@@ -32,7 +32,7 @@ TracedRational{T}(c::AbstractChar) where {T} = TracedRational{T}(codepoint(c))
 TracedRational(c::AbstractChar) = TracedRational(codepoint(c))
 
 function TracedRational{T}(::Complex) where {T}
-    throw("Currently we dont support complex numbers in rationals")
+    return throw("Currently we dont support complex numbers in rationals")
 end
 TracedRational(::Complex) = throw("Currently we dont support complex numbers in rationals")
 
@@ -60,7 +60,7 @@ end
 # Basic properties
 Base.numerator(x::TracedRational) = x.num
 Base.denominator(x::TracedRational) = x.den
-function Base.denominator(::TracedRNumber{T}) where {T<:Integer}
+function Base.denominator(::TracedRInteger{T}) where {T}
     return Reactant.promote_to(TracedRNumber, one(T))
 end
 
@@ -124,11 +124,11 @@ for op in (:+, :-, :rem, :mod)
             return TracedRational($(op)(x.num * yd, y.num * xd), x.den * yd)
         end
 
-        function Base.$(op)(x::TracedRational, y::Union{TracedRNumber{<:Integer},Integer})
+        function Base.$(op)(x::TracedRational, y::Integer)
             return TracedRational($(op)(x.num, x.den * y), x.den)
         end
 
-        function Base.$(op)(x::Union{TracedRNumber{<:Integer},Integer}, y::TracedRational)
+        function Base.$(op)(x::Integer, y::TracedRational)
             return TracedRational($(op)(x * y.den, y.num), y.den)
         end
     end
@@ -140,12 +140,12 @@ function Base.:*(x::TracedRational, y::TracedRational)
     return TracedRational(xn * yn, xd * yd)
 end
 
-function Base.:*(x::TracedRational, y::Union{TracedRNumber{<:Integer},Integer})
+function Base.:*(x::TracedRational, y::Integer)
     xd, yn = Base.divgcd(promote(x.den, y)...)
     return TracedRational(x.num * yn, xd)
 end
 
-function Base.:*(x::Union{TracedRNumber{<:Integer},Integer}, y::TracedRational)
+function Base.:*(x::Integer, y::TracedRational)
     xn, yd = Base.divgcd(promote(x, y.den)...)
     return TracedRational(xn * y.num, yd)
 end
@@ -172,13 +172,10 @@ Base.one(::Type{TracedRational{T}}) where {T} = TracedRational(one(T), one(T))
 Base.zero(::TracedRational{T}) where {T} = zero(T)
 Base.one(::TracedRational{T}) where {T} = one(T)
 
-# Rational construction with // operator
-for (T1, T2) in Iterators.product(
-    (Integer, TracedRNumber{<:Integer}, AbstractConcreteNumber{<:Integer}),
-    (Integer, TracedRNumber{<:Integer}, AbstractConcreteNumber{<:Integer}),
-)
-    T1 == T2 && T1 == Integer && continue
-
+# Rational construction with the // operator. `RInteger <: Integer` covers
+# traced and concrete Reactant integers, so three methods suffice and dominate
+# `//(::Integer, ::Integer)` in Base.
+for (T1, T2) in ((RInteger, RInteger), (RInteger, Integer), (Integer, RInteger))
     @eval function Base.://(num::$(T1), den::$(T2))
         return TracedRational(
             promote(
@@ -192,12 +189,12 @@ end
 Base.://(num::TracedRational, den::Rational) = num//TracedRational(den)
 Base.://(num::Rational, den::TracedRational) = TracedRational(num)//den
 
-function Base.://(x::TracedRational, y::Union{TracedRNumber{<:Integer},Integer})
+function Base.://(x::TracedRational, y::Integer)
     xn, yn = Base.divgcd(promote(x.num, y)...)
     return checked_den(xn, x.den * yn)
 end
 
-function Base.://(x::Union{TracedRNumber{<:Integer},Integer}, y::TracedRational)
+function Base.://(x::Integer, y::TracedRational)
     xn, yn = Base.divgcd(promote(x, y.num)...)
     return checked_den(xn * y.den, yn)
 end
@@ -208,8 +205,8 @@ function Base.://(x::TracedRational, y::TracedRational)
     return checked_den(xn * yd, xd * yn)
 end
 
-Base.:/(x::TracedRational, y::Union{TracedRNumber{<:Integer},Integer}) = x//y
-Base.:/(x::Union{TracedRNumber{<:Integer},Integer}, y::TracedRational) = x//y
+Base.:/(x::TracedRational, y::Integer) = x//y
+Base.:/(x::Integer, y::TracedRational) = x//y
 Base.:/(x::TracedRational, y::TracedRational) = x//y
 
 end
diff --git a/src/TracedUtils.jl b/src/TracedUtils.jl
index a1bb430e32..a912f57664 100644
--- a/src/TracedUtils.jl
+++ b/src/TracedUtils.jl
@@ -81,7 +81,9 @@ function ReactantCore.materialize_traced_array(
     return permutedims(materialize_traced_array(parent(x)), perm)
 end
 
-function ReactantCore.materialize_traced_array(x::AbstractArray{TracedRNumber{T}}) where {T}
+function ReactantCore.materialize_traced_array(
+    x::AbstractArray{<:TracedRNumber{T}}
+) where {T}
     as = Reactant.aos_to_soa(x)
     if as === x
         as = x[axes(x)...]
@@ -90,17 +92,17 @@ function ReactantCore.materialize_traced_array(x::AbstractArray{TracedRNumber{T}
 end
 
 get_mlir_data(x::TracedRNumber) = x.mlir_data
-set_mlir_data!(x::TracedRNumber, data) = (x.mlir_data = data; return x)
+set_mlir_data!(x::TracedRNumber, data) = (x.mlir_data=data; return x)
 get_paths(x::TracedRNumber) = x.paths
-set_paths!(x::TracedRNumber, paths) = (x.paths = paths; return x)
+set_paths!(x::TracedRNumber, paths) = (x.paths=paths; return x)
 
 get_mlir_data(x::TracedRArray) = x.mlir_data
 get_mlir_data(x::AnyTracedRArray) = get_mlir_data(materialize_traced_array(x))
 get_paths(x::TracedRArray) = x.paths
-set_paths!(x::TracedRArray, paths) = (x.paths = paths; return x)
+set_paths!(x::TracedRArray, paths) = (x.paths=paths; return x)
 
 get_paths(x::MissingTracedValue) = x.paths
-set_paths!(x::MissingTracedValue, paths) = (x.paths = paths; return x)
+set_paths!(x::MissingTracedValue, paths) = (x.paths=paths; return x)
 
 function set_mlir_data!(x::TracedRArray, data)
     x.mlir_data = data
@@ -136,7 +138,7 @@ function get_ancestor_and_indices(
             bcasted_idxs = @opcall broadcast_in_dim(
                 idx, ndims(idx) == 0 ? Int64[] : Int64[i], flattened_size
             )
-            Base.stride(x, i) .* (bcasted_idxs .- 1)
+            return Base.stride(x, i) .* (bcasted_idxs .- 1)
         end
         linear_indices = linear_indices .+ 1
         parent_linear_indices_all = collect(LinearIndices(size(parent(x))))
@@ -220,7 +222,7 @@ __to_cartesian_index(x::NTuple{N,<:TracedRNumber}) where {N} = vcat(x...)
 function _get_ancestor_and_indices_linear(x::AnyTracedRArray, indices::AbstractArray)
     pidxs = parentindices(x)
     parent_indices = map(CartesianIndices(x)[indices]) do idx
-        __to_cartesian_index(Base.reindex(pidxs, (idx.I...,)))
+        return __to_cartesian_index(Base.reindex(pidxs, (idx.I...,)))
     end
     return get_ancestor_and_indices(parent(x), reduce(vcat, parent_indices))
 end
@@ -1094,7 +1096,7 @@ function get_idx(x, prefix::Symbol)
             return path
         end
     end
-    throw(AssertionError("No path found for $x"))
+    return throw(AssertionError("No path found for $x"))
 end
 
 function get_argidx(x, prefix::Symbol)
@@ -1288,12 +1290,11 @@ function elem_apply(f, args::Vararg{Any,Nargs}) where {Nargs}
         push_val!(batch_inputs, ogarg, path[3:end])
 
         if ogarg isa Base.RefValue
-            batch_inputs[end] =
-                (@opcall broadcast_in_dim(
-                    TracedRArray(batch_inputs[end]),
-                    Int64[],
-                    collect(Int64, input_shapes[1]),
-                )).mlir_data
+            batch_inputs[end] = (@opcall broadcast_in_dim(
+                TracedRArray(batch_inputs[end]),
+                Int64[],
+                collect(Int64, input_shapes[1]),
+            )).mlir_data
         end
     end
 

From cd98dd367b4f6b302545472507b5f0f43df827a8 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Tue, 7 Jul 2026 20:13:16 +0200
Subject: [PATCH 07/30] Adapt concrete number ops, codegen, and Enzyme
 integration

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.
---
 src/ConcreteRArray.jl   | 117 +++++++++++++++++++++---------
 src/Enzyme.jl           |  38 ++++++++--
 src/compiler/Codegen.jl | 153 ++++++++++++++++++++++------------------
 3 files changed, 199 insertions(+), 109 deletions(-)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index cc341ded46..4ca5afe420 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -53,30 +53,37 @@ function Base.deepcopy_internal(
 end
 
 Base.size(::AbstractConcreteNumber) = ()
-Base.real(x::AbstractConcreteNumber{<:Real}) = x
 function Base.rtoldefault(T::Type{<:AbstractConcreteNumber})
     return T(Base.rtoldefault(unwrapped_eltype(T)))
 end
+function Base.rtoldefault(T::Type{<:AbstractConcreteFloat})
+    return T(Base.rtoldefault(unwrapped_eltype(T)))
+end
 
 Base.strides(x::AbstractConcreteArray) = Base.size_to_strides(1, size(x)...)
 
-Base.OneTo(x::AbstractConcreteNumber{<:Integer}) = Base.OneTo(to_number(x))
+Base.OneTo(x::AbstractConcreteInteger) = Base.OneTo(to_number(x))
 
 @static if isdefined(Base, :unchecked_oneto)
-    function Base.unchecked_oneto(x::AbstractConcreteNumber{<:Integer})
+    function Base.unchecked_oneto(x::AbstractConcreteInteger)
         return Base.unchecked_oneto(to_number(x))
     end
 end
 
 # Ensure the device and client are the same as the input
-for numType in (:ConcretePJRTNumber, :ConcreteIFRTNumber)
-    @eval function Base.float(x::$(numType){T}) where {T}
-        return $(numType)(
+for (numType, familyType) in Iterators.product(
+    (:Integer, :Float, :Complex), (:ConcretePJRTNumber, :ConcreteIFRTNumber)
+)
+    leafType = Symbol(replace(string(familyType), "Number" => string(numType)))
+    @eval function Base.float(x::$(leafType){T}) where {T}
+        return $(familyType)(
             float(T)(to_number(x)); client=XLA.client(x), device=XLA.device(x), x.sharding
         )
     end
 end
 
+Base.decompose(x::AbstractConcreteFloat) = Base.decompose(to_number(x))
+
 # written like this to avoid ambiguity errors
 for T in Base.uniontypes(ReactantPrimitive)
     @eval (::Type{$(T)})(x::AbstractConcreteNumber) = convert($T, x)
@@ -202,28 +209,63 @@ function Base.convert(
     return to_number(x)
 end
 
-for jlop in (:(Base.abs),), T in (AbstractConcreteNumber,)
+for jlop in (:(Base.abs),),
+    T in (AbstractConcreteInteger, AbstractConcreteFloat, AbstractConcreteComplex)
+
     @eval $(jlop)(x::$(T)) = $(jlop)(to_number(x))
 end
 
+# The concrete number methods are defined per numeric kind: a single method on
+# the `AbstractConcreteNumber` union would be ambiguous with Base's methods on
+# `Real`/`AbstractFloat`/`Integer`. The extra methods for `Rational`, `BigInt`,
+# `BigFloat`, `AbstractIrrational`, `AbstractFloat`, and `Complex` resolve
+# ambiguities with Base's specialized methods, and the traced-number methods
+# preserve the deliberate error on mixing concrete and traced numbers.
+const CONCRETE_NUMBER_KINDS = (
+    AbstractConcreteInteger, AbstractConcreteFloat, AbstractConcreteComplex
+)
+const TRACED_NUMBER_KINDS = (TracedRInteger, TracedRFloat, TracedRComplex)
+
 for jlop in (
-        :(Base.isless),
-        :(Base.:+),
-        :(Base.:-),
-        :(Base.:*),
-        :(Base.:/),
-        :(Base.:^),
-        :(Base.:(==)),
-    ),
-    T in (AbstractConcreteNumber, AbstractConcreteArray{<:Number,0})
+    :(Base.isless),
+    :(Base.:+),
+    :(Base.:-),
+    :(Base.:*),
+    :(Base.:/),
+    :(Base.:^),
+    :(Base.:(==)),
+)
+    for T1 in CONCRETE_NUMBER_KINDS
+        @eval begin
+            $(jlop)(x::$(T1), y::Number) = $(jlop)(to_number(x), y)
+            $(jlop)(x::Number, y::$(T1)) = $(jlop)(x, to_number(y))
+        end
+        for T2 in CONCRETE_NUMBER_KINDS
+            @eval $(jlop)(x::$(T1), y::$(T2)) = $(jlop)(to_number(x), to_number(y))
+        end
+        for X in
+            (:Rational, :BigInt, :BigFloat, :AbstractIrrational, :AbstractFloat, :Complex)
+            @eval begin
+                $(jlop)(x::$(T1), y::$(X)) = $(jlop)(to_number(x), y)
+                $(jlop)(x::$(X), y::$(T1)) = $(jlop)(x, to_number(y))
+            end
+        end
+        for T2 in TRACED_NUMBER_KINDS
+            @eval begin
+                $(jlop)(x::$(T1), y::$(T2)) = throw(MethodError($(jlop), (x, y)))
+                $(jlop)(x::$(T2), y::$(T1)) = throw(MethodError($(jlop), (x, y)))
+            end
+        end
+    end
 
+    T = AbstractConcreteArray{<:Number,0}
     @eval begin
         $(jlop)(x::$(T), y::$(T)) = $(jlop)(to_number(x), to_number(y))
         $(jlop)(x::$(T), y::Number) = $(jlop)(to_number(x), y)
         $(jlop)(x::Number, y::$(T)) = $(jlop)(x, to_number(y))
 
-        $(jlop)(x::$(T), y::TracedRNumber) = throw(MethodError(jlop, (x, y)))
-        $(jlop)(x::TracedRNumber, y::$(T)) = throw(MethodError(jlop, (x, y)))
+        $(jlop)(x::$(T), y::TracedRNumber) = throw(MethodError($(jlop), (x, y)))
+        $(jlop)(x::TracedRNumber, y::$(T)) = throw(MethodError($(jlop), (x, y)))
     end
 end
 
@@ -249,13 +291,14 @@ for T in (AbstractConcreteNumber, AbstractConcreteArray{<:Number,0})
         div(to_number(x), y, r)
 end
 
-for T in (Integer, Rational)
-    @eval Base.:^(x::AbstractConcreteNumber, y::$(T)) = ^(to_number(x), y)
+# `Rational` exponents are covered by the disambiguation grid above.
+for CT in CONCRETE_NUMBER_KINDS
+    @eval Base.:^(x::$(CT), y::Integer) = ^(to_number(x), y)
 end
-Base.:^(::Irrational{:ℯ}, x::AbstractConcreteNumber) = exp(x)
+Base.:^(::Irrational{:ℯ}, x::AbstractConcreteNumber) = exp(to_number(x))
 
 for jlop in (:(Base.isnan), :(Base.isfinite)),
-    T in (AbstractConcreteNumber, AbstractConcreteArray{<:Any,0})
+    T in (CONCRETE_NUMBER_KINDS..., AbstractConcreteArray{<:Any,0})
 
     @eval $(jlop)(x::$(T)) = $(jlop)(to_number(x))
 end
@@ -297,7 +340,7 @@ for (T1, T2) in (
     end
 end
 
-function Base.show(io::IO, X::Union{ConcretePJRTScalar,ConcreteIFRTScalar})
+function show_concrete_scalar(io::IO, X)
     if isempty(X)
         print(io, "<Empty Buffer eltype $(eltype(X)) of size $(size(X))>")
         return nothing
@@ -308,6 +351,10 @@ function Base.show(io::IO, X::Union{ConcretePJRTScalar,ConcreteIFRTScalar})
     return nothing
 end
 
+for T in (CONCRETE_NUMBER_KINDS..., ConcretePJRTArray{<:Any,0}, ConcreteIFRTArray{<:Any,0})
+    @eval Base.show(io::IO, X::$(T)) = show_concrete_scalar(io, X)
+end
+
 function Base.print_array(io::IO, X::Union{AnyConcretePJRTArray,AnyConcreteIFRTArray})
     if isempty(X)
         print(io, "<Empty Buffer eltype $(eltype(X)) of size $(size(X))>")
@@ -477,7 +524,7 @@ function Base.similar(
     @assert length(device_to_array_slices) == D
     sdata = ntuple(Val(D)) do i
         Base.@_inline_meta
-        similar(a.data[i], S, Dims(length.(device_to_array_slices[i])))
+        return similar(a.data[i], S, Dims(length.(device_to_array_slices[i])))
     end
     return ConcretePJRTArray{S,length(dims),D}(sdata, dims, a.sharding)
 end
@@ -542,7 +589,7 @@ function Base.copy(bc::Base.Broadcast.Broadcasted{Broadcast.ArrayStyle{ConcreteP
     if all(buffer_on_cpu, bc.args) && all(
         x ->
             !(x isa ConcretePJRTArray) ||
-                (x isa ConcretePJRTArray && !Sharding.is_sharded(x)),
+            (x isa ConcretePJRTArray && !Sharding.is_sharded(x)),
         bc.args,
     )
         ElType = Base.Broadcast.combine_eltypes(bc.f, bc.args)
@@ -787,7 +834,7 @@ end
 function Base.fill!(
     a::ConcretePJRTArray{<:TracedRNumber{T},N}, val::TracedRNumber{T2}
 ) where {T,T2,N}
-    throw(MethodError(fill!, (a, val)))
+    return throw(MethodError(fill!, (a, val)))
 end
 
 function Base.fill!(a::ConcretePJRTArray{T,N}, val) where {T,N}
@@ -821,7 +868,7 @@ end
 function Base.fill!(
     a::ConcreteIFRTArray{<:TracedRNumber{T},N}, val::TracedRNumber{T2}
 ) where {T,T2,N}
-    throw(MethodError(fill!, (a, val)))
+    return throw(MethodError(fill!, (a, val)))
 end
 
 function Base.fill!(x::UnionAnyConcreteRArray, val)
@@ -833,19 +880,19 @@ end
 function Base.fill!(
     x::Union{ConcreteIFRTArray{<:TracedRNumber},ConcretePJRTArray{<:TracedRNumber}}, val
 )
-    throw(MethodError(fill!, (x, val)))
+    return throw(MethodError(fill!, (x, val)))
 end
 function Base.fill!(
     x::Union{ConcreteIFRTArray{<:TracedRNumber},ConcretePJRTArray{<:TracedRNumber}},
     val::TracedRNumber,
 )
-    throw(MethodError(fill!, (x, val)))
+    return throw(MethodError(fill!, (x, val)))
 end
 function Base.fill!(
     x::Union{AnyConcreteIFRTArray{<:TracedRNumber},AnyConcretePJRTArray{<:TracedRNumber}},
     val,
 )
-    throw(MethodError(fill!, (x, val)))
+    return throw(MethodError(fill!, (x, val)))
 end
 
 function mymapreducedim!(f, op, R, A)
@@ -901,8 +948,12 @@ for T in (Number, Integer)
     end
 end
 
-Base.isinf(x::ConcreteRNumber{T}) where {T} = Base.isinf(convert(T, x))
-Base.round(x::ConcreteRNumber{T}) where {T} = Base.round(convert(T, x))
+for CT in CONCRETE_NUMBER_KINDS
+    @eval begin
+        Base.isinf(x::$(CT){T}) where {T} = Base.isinf(convert(T, x))
+        Base.round(x::$(CT){T}) where {T} = Base.round(convert(T, x))
+    end
+end
 
 Base._parentsmatch(A::ConcreteIFRTArray, B::ConcreteIFRTArray) = A === B
 Base._parentsmatch(A::ConcretePJRTArray, B::ConcretePJRTArray) = A === B
@@ -933,4 +984,4 @@ for srcStyle in (IndexStyle, IndexCartesian),
     end
 end
 
-Base.to_index(x::AbstractConcreteNumber) = to_number(x)
+Base.to_index(x::AbstractConcreteInteger) = to_number(x)
diff --git a/src/Enzyme.jl b/src/Enzyme.jl
index 26995f32b3..987f68dc02 100644
--- a/src/Enzyme.jl
+++ b/src/Enzyme.jl
@@ -61,6 +61,22 @@ end
     return zero(Core.Typeof(x))
 end
 
+# Per-kind methods: `RFloat <: AbstractFloat`, so a single `RNumber` method
+# would be ambiguous with Enzyme's methods on `AbstractFloat`.
+for RT in (:RInteger, :RFloat, :RComplex)
+    @eval @inline Enzyme.make_zero(x::$RT) = zero(Core.Typeof(x))
+end
+
+# Reactant numbers are mutable wrappers around device/MLIR values and must
+# stay `Duplicated` under autodiff even though they subtype `AbstractFloat`.
+for RT in (:RArray, :RNumber, :RInteger, :RFloat, :RComplex)
+    @eval @inline function Enzyme.guess_activity(
+        ::Type{T}, mode::Enzyme.Mode
+    ) where {T<:$RT}
+        return Enzyme.Duplicated{T}
+    end
+end
+
 @inline function Enzyme.make_zero(x::RArray{FT,N})::RArray{FT,N} where {FT<:AbstractFloat,N}
     return Base.zero(x)
 end
@@ -97,9 +113,9 @@ end
 @register_make_zero_inplace(Enzyme.make_zero!)
 @register_make_zero_inplace(Enzyme.remake_zero!)
 
-function Enzyme.make_zero(
-    ::Type{RT}, seen::IdDict, prev::RT, ::Val{copy_if_inactive}=Val(false)
-)::RT where {copy_if_inactive,RT<:Union{RArray,RNumber}}
+function make_zero_impl(
+    ::Type{RT}, seen::IdDict, prev::RT, ::Val{copy_if_inactive}
+)::RT where {copy_if_inactive,RT}
     if haskey(seen, prev)
         return seen[prev]
     end
@@ -111,6 +127,14 @@ function Enzyme.make_zero(
     return res
 end
 
+for T in (:(Union{RArray,RNumber}), :RInteger, :RFloat, :RComplex)
+    @eval function Enzyme.make_zero(
+        ::Type{RT}, seen::IdDict, prev::RT, ::Val{copy_if_inactive}=Val(false)
+    )::RT where {copy_if_inactive,RT<:$T}
+        return make_zero_impl(RT, seen, prev, Val(copy_if_inactive))
+    end
+end
+
 function Enzyme.onehot(x::TracedRArray{T,N}) where {T,N}
     onehot_matrix = promote_to(TracedRArray{T,2}, LinearAlgebra.I(length(x)))
     return Tuple(
@@ -144,7 +168,7 @@ function EnzymeRules.augmented_primal(
 ) where {RT}
     primargs = ntuple(Val(length(args))) do i
         Base.@_inline_meta
-        args[i].val
+        return args[i].val
     end
 
     primal = if EnzymeCore.needs_primal(config)
@@ -164,7 +188,7 @@ function EnzymeRules.augmented_primal(
         else
             ntuple(Val(EnzymeRules.width(config))) do i
                 Base.@_inline_meta
-                ConcretePJRTArray(
+                return ConcretePJRTArray(
                     zeros(T.val, primargs...);
                     client=XLA.client(uval.val),
                     device=XLA.device(uval.val),
@@ -194,7 +218,7 @@ function EnzymeRules.reverse(
 ) where {RT,N}
     ntuple(Val(N + 2)) do i
         Base.@_inline_meta
-        nothing
+        return nothing
     end
 end
 
@@ -497,7 +521,7 @@ function overload_autodiff(
         else
             ntuple(Val(width)) do i
                 Base.@_inline_meta
-                deepcopy(result)
+                return deepcopy(result)
             end
         end
     else
diff --git a/src/compiler/Codegen.jl b/src/compiler/Codegen.jl
index 0c092843ce..6213e851f1 100644
--- a/src/compiler/Codegen.jl
+++ b/src/compiler/Codegen.jl
@@ -6,8 +6,14 @@ using ..Reactant:
     TracedRNumber,
     ConcretePJRTArray,
     ConcretePJRTNumber,
+    ConcretePJRTInteger,
+    ConcretePJRTFloat,
+    ConcretePJRTComplex,
     ConcreteIFRTArray,
     ConcreteIFRTNumber,
+    ConcreteIFRTInteger,
+    ConcreteIFRTFloat,
+    ConcreteIFRTComplex,
     AbstractConcreteArray,
     AbstractConcreteNumber,
     ancestor
@@ -241,12 +247,9 @@ Base.@nospecializeinfer function create_result(
 
         result = Expr(:new, T, elems...)
 
-        push!(
-            resultgen_code,
-            quote
-                $sym = $result
-            end,
-        )
+        push!(resultgen_code, quote
+            $sym = $result
+        end)
         result_cache[tocopy] = sym
     end
 
@@ -272,19 +275,16 @@ Base.@nospecializeinfer function create_result(
 
         result = Meta.quot(tocopy)
 
-        push!(
-            resultgen_code,
-            quote
-                $sym = $result
-            end,
-        )
+        push!(resultgen_code, quote
+            $sym = $result
+        end)
         result_cache[tocopy] = sym
     end
 
     return result_cache[tocopy]
 end
 
-function create_result(
+function create_pjrt_number_result(
     tocopy::ConcretePJRTNumber{T,D},
     @nospecialize(path::Tuple),
     result_stores,
@@ -314,19 +314,16 @@ function create_result(
         else
             result = :(ConcretePJRTNumber{$T}($restore))
         end
-        push!(
-            resultgen_code,
-            quote
-                $sym = $result
-            end,
-        )
+        push!(resultgen_code, quote
+            $sym = $result
+        end)
         result_cache[tocopy] = sym
     end
 
     return result_cache[tocopy]
 end
 
-function create_result(
+function create_ifrt_number_result(
     tocopy::ConcreteIFRTNumber{T},
     @nospecialize(path::Tuple),
     result_stores,
@@ -356,18 +353,57 @@ function create_result(
         else
             result = :(ConcreteIFRTNumber{$T}($restore))
         end
-        push!(
-            resultgen_code,
-            quote
-                $sym = $result
-            end,
-        )
+        push!(resultgen_code, quote
+            $sym = $result
+        end)
         result_cache[tocopy] = sym
     end
 
     return result_cache[tocopy]
 end
 
+# The per-kind methods are needed (rather than a single `ConcretePJRTNumber`
+# method) so that integer/float scalars are not ambiguous with the
+# `create_result(::Union{Integer,AbstractFloat,...})` method below.
+for (NumT, worker) in (
+    (:ConcretePJRTNumber, :create_pjrt_number_result),
+    (:ConcretePJRTInteger, :create_pjrt_number_result),
+    (:ConcretePJRTFloat, :create_pjrt_number_result),
+    (:ConcretePJRTComplex, :create_pjrt_number_result),
+    (:ConcreteIFRTNumber, :create_ifrt_number_result),
+    (:ConcreteIFRTInteger, :create_ifrt_number_result),
+    (:ConcreteIFRTFloat, :create_ifrt_number_result),
+    (:ConcreteIFRTComplex, :create_ifrt_number_result),
+)
+    @eval function create_result(
+        tocopy::$NumT,
+        @nospecialize(path::Tuple),
+        result_stores,
+        path_to_shard_info,
+        to_unreshard_results,
+        unresharded_code::Vector{Expr},
+        unresharded_arrays_cache,
+        used_shardinfo,
+        result_cache,
+        var_idx,
+        resultgen_code,
+    )
+        return $worker(
+            tocopy,
+            path,
+            result_stores,
+            path_to_shard_info,
+            to_unreshard_results,
+            unresharded_code,
+            unresharded_arrays_cache,
+            used_shardinfo,
+            result_cache,
+            var_idx,
+            resultgen_code,
+        )
+    end
+end
+
 function create_result(
     tocopy::ConcretePJRTArray{T,N,D},
     @nospecialize(path::Tuple),
@@ -399,12 +435,9 @@ function create_result(
         else
             result = :(ConcretePJRTArray{$T,$N}($restore, $(tocopy.shape)))
         end
-        push!(
-            resultgen_code,
-            quote
-                $sym = $result
-            end,
-        )
+        push!(resultgen_code, quote
+            $sym = $result
+        end)
         result_cache[tocopy] = sym
     end
 
@@ -458,12 +491,9 @@ function create_result(
         else
             result = :(ConcreteIFRTArray{$T,$N}($(restore), $(tocopy.shape)))
         end
-        push!(
-            resultgen_code,
-            quote
-                $sym = $result
-            end,
-        )
+        push!(resultgen_code, quote
+            $sym = $result
+        end)
         result_cache[tocopy] = sym
     end
 
@@ -526,23 +556,17 @@ function create_result(
         sym = Symbol("result", var_idx[])
         var_idx[] += 1
 
-        push!(
-            resultgen_code,
-            quote
-                $sym = $(Array{T,N})(undef, $(size(tocopy)...,))
-            end,
-        )
+        push!(resultgen_code, quote
+            $sym = $(Array{T,N})(undef, $(size(tocopy)...,))
+        end)
 
         result_cache[tocopy] = sym
 
         for (i, v) in enumerate(tocopy)
             subexpr = create_result(v, append_path(path, i), args...)
-            push!(
-                resultgen_code,
-                quote
-                    @inbounds $sym[$i] = $subexpr
-                end,
-            )
+            push!(resultgen_code, quote
+                @inbounds $sym[$i] = $subexpr
+            end)
         end
     end
 
@@ -640,12 +664,9 @@ function create_result(
         sym = Symbol("result", var_idx[])
         var_idx[] += 1
 
-        push!(
-            resultgen_code,
-            quote
-                $sym = $D()
-            end,
-        )
+        push!(resultgen_code, quote
+            $sym = $D()
+        end)
 
         result_cache[tocopy] = sym
 
@@ -654,12 +675,9 @@ function create_result(
             # symbol keys must be quoted in generated code; otherwise
             # they are interpreted as variable references
             k_expr = k isa Symbol ? QuoteNode(k) : k
-            push!(
-                resultgen_code,
-                quote
-                    @inbounds $sym[$k_expr] = $subexpr
-                end,
-            )
+            push!(resultgen_code, quote
+                @inbounds $sym[$k_expr] = $subexpr
+            end)
         end
     end
 
@@ -1213,12 +1231,9 @@ function codegen_unflatten!(
             sym = Symbol("result", var_idx[])
             var_idx[] += 1
 
-            push!(
-                resultgen_code,
-                quote
-                    $sym = $argres.data
-                end,
-            )
+            push!(resultgen_code, quote
+                $sym = $argres.data
+            end)
 
             result_stores[path] = sym
         end
@@ -1345,7 +1360,7 @@ function codegen_xla_call(
     base_symbol_name = is_sharded ? Symbol(:result_buffer_m, ndevices, :_) : :result_buffer_
     concretized_res_names = Symbol[Symbol(base_symbol_name, i) for i in 1:nresults]
     concretized_res_code = map(enumerate(concretized_res_names)) do (i, varname)
-        :($varname = linearized_results[$i])
+        return :($varname = linearized_results[$i])
     end
 
     xla_call_code = if nresults == 0 && is_pure

From 8e15817111844b19918a474688f3e7249f7647b9 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Tue, 7 Jul 2026 20:13:16 +0200
Subject: [PATCH 08/30] Adapt stdlib overloads and package extensions to the
 split number types

Created by generative AI.
---
 ext/ReactantAbstractFFTsExt.jl          |  8 ++++----
 ext/ReactantCUDAExt.jl                  | 10 ++++++----
 ext/ReactantKernelAbstractionsExt.jl    |  3 +--
 ext/ReactantMPIExt/Ops.jl               |  4 ++--
 ext/ReactantOneHotArraysExt.jl          | 16 ++++++++++------
 ext/ReactantSparseArraysExt/ReadOnly.jl |  2 +-
 src/stdlibs/LinearAlgebra.jl            | 14 ++++++++------
 7 files changed, 32 insertions(+), 25 deletions(-)

diff --git a/ext/ReactantAbstractFFTsExt.jl b/ext/ReactantAbstractFFTsExt.jl
index 3b5aeff7ee..3cde0dd65a 100644
--- a/ext/ReactantAbstractFFTsExt.jl
+++ b/ext/ReactantAbstractFFTsExt.jl
@@ -121,7 +121,7 @@ for op in (:rfft, :fft, :ifft)
         if __is_valid_stablehlo_fft_dims(dims, ndims(x))
             return @opcall fft(
                 TracedUtils.materialize_traced_array(x);
-                type=$(uppercase(string(op))),
+                type=($(uppercase(string(op)))),
                 length=fft_lengths,
             )
         end
@@ -130,7 +130,7 @@ for op in (:rfft, :fft, :ifft)
             @opcall(
                 fft(
                     TracedUtils.materialize_traced_array(permutedims(x, perm));
-                    type=$(uppercase(string(op))),
+                    type=($(uppercase(string(op)))),
                     length=fft_lengths,
                 )
             ),
@@ -219,7 +219,7 @@ for op in (:irfft,)
         if __is_valid_stablehlo_fft_dims(dims, ndims(x))
             return @opcall fft(
                 TracedUtils.materialize_traced_array(x);
-                type=$(uppercase(string(op))),
+                type=($(uppercase(string(op)))),
                 length=fft_lengths,
             )
         end
@@ -229,7 +229,7 @@ for op in (:irfft,)
             @opcall(
                 fft(
                     TracedUtils.materialize_traced_array(permutedims(x, perm));
-                    type=$(uppercase(string(op))),
+                    type=($(uppercase(string(op)))),
                     length=fft_lengths,
                 )
             ),
diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl
index 35ea0098a2..dabadd2bfb 100644
--- a/ext/ReactantCUDAExt.jl
+++ b/ext/ReactantCUDAExt.jl
@@ -808,10 +808,12 @@ function compile(job)
             prevdl = String(prevdlattr)
             @assert prevdl == dl "data layout mismatch, tried compiling cuda kernels for different target machines?"
         else
-            MLIR.IR.setattr!(MLIR.IR.Operation(cur_module), dl_attr_name, MLIR.IR.Attribute(dl))
+            MLIR.IR.setattr!(
+                MLIR.IR.Operation(cur_module), dl_attr_name, MLIR.IR.Attribute(dl)
+            )
         end
 
-        String(Reactant.TracedUtils.get_attribute_by_name(linkRes, "sym_name"))
+        return String(Reactant.TracedUtils.get_attribute_by_name(linkRes, "sym_name"))
     end
 
     return LLVMFunc{job.source.specTypes.parameters[1],job.source.specTypes}(nothing, entry)
@@ -1061,7 +1063,7 @@ Reactant.@reactant_overlay function (func::LLVMFunc{F,tt})(
     has_cast_float_type =
         bfloat16_compile_type !== BFloat16 && any(values(seen)) do arg
             (arg isa TracedRArray || arg isa TracedRNumber) &&
-                Reactant.unwrapped_eltype(typeof(arg)) === BFloat16
+            Reactant.unwrapped_eltype(typeof(arg)) === BFloat16
         end
 
     wrapper_tys = MLIR.IR.Type[]
@@ -1516,7 +1518,7 @@ const ReactantCUDAJob = GPUCompiler.CompilerJob{
 GPUCompiler.can_vectorize(job::ReactantCUDAJob) = !job.config.params.raising
 function GPUCompiler.optimization_options(job::ReactantCUDAJob)
     raising = job.config.params.raising
-    return (; instcombine=!raising, fastmath=!raising, aggressiveinstcombine=!raising)
+    return (; instcombine=(!raising), fastmath=(!raising), aggressiveinstcombine=(!raising))
 end
 
 function GPUCompiler.method_table(@nospecialize(job::ReactantCUDAJob))
diff --git a/ext/ReactantKernelAbstractionsExt.jl b/ext/ReactantKernelAbstractionsExt.jl
index 954b153c92..fa9ccaccf0 100644
--- a/ext/ReactantKernelAbstractionsExt.jl
+++ b/ext/ReactantKernelAbstractionsExt.jl
@@ -112,9 +112,8 @@ end
 @static if VERSION < v"1.12-"
     Reactant.@reactant_overlay Base.@nospecializeinfer @noinline function (
         obj::KA.Kernel{ReactantBackend}
-    )(
-        @nospecialize args...; ndrange=nothing, workgroupsize=nothing
     )
+        (@nospecialize args...; ndrange=nothing, workgroupsize=nothing)
         return Reactant.call_with_reactant(
             Reactant.ka_with_reactant, ndrange, workgroupsize, obj, args...
         )
diff --git a/ext/ReactantMPIExt/Ops.jl b/ext/ReactantMPIExt/Ops.jl
index 88eb0db359..4bf0e04ce4 100644
--- a/ext/ReactantMPIExt/Ops.jl
+++ b/ext/ReactantMPIExt/Ops.jl
@@ -225,7 +225,7 @@ const MPI_OP_MAP = Dict(
 
 function get_mpi_op_enum(op)
     return get(MPI_OP_MAP, op.val) do
-        throw(ArgumentError("Unknown MPI op `$op`"))
+        return throw(ArgumentError("Unknown MPI op `$op`"))
     end
 end
 
@@ -261,7 +261,7 @@ const MPI_DATATYPE_MAP = Dict(
 
 function get_mpi_datatype_enum(datatype)
     return get(MPI_DATATYPE_MAP, datatype.val) do
-        throw(ArgumentError("Unknown MPI datatype `$datatype`"))
+        return throw(ArgumentError("Unknown MPI datatype `$datatype`"))
     end
 end
 
diff --git a/ext/ReactantOneHotArraysExt.jl b/ext/ReactantOneHotArraysExt.jl
index 4a7e3ddfdf..94290f2300 100644
--- a/ext/ReactantOneHotArraysExt.jl
+++ b/ext/ReactantOneHotArraysExt.jl
@@ -7,11 +7,15 @@ using ReactantCore: ReactantCore
 using Reactant.Ops: @opcall
 
 __compatible_eltype(::Type{T}, ::Type{U}) where {T,U} = T
-function __compatible_eltype(::Type{TracedRNumber{T}}, ::Type{TracedRNumber{U}}) where {T,U}
-    return TracedRNumber{T}
+function __compatible_eltype(
+    ::Type{<:TracedRNumber{T}}, ::Type{<:TracedRNumber{U}}
+) where {T,U}
+    return Reactant.traced_number_type(T)
+end
+__compatible_eltype(::Type{<:TracedRNumber{T}}, ::Type{U}) where {T,U} = T
+function __compatible_eltype(::Type{T}, ::Type{<:TracedRNumber{U}}) where {T,U}
+    return Reactant.traced_number_type(T)
 end
-__compatible_eltype(::Type{TracedRNumber{T}}, ::Type{U}) where {T,U} = T
-__compatible_eltype(::Type{T}, ::Type{TracedRNumber{U}}) where {T,U} = TracedRNumber{T}
 
 function Reactant.traced_type_inner(
     @nospecialize(_::Type{OneHotArray{T,N,Np1,I}}),
@@ -60,7 +64,7 @@ function OneHotArrays.onehotbatch(data::AnyTracedRArray{<:Any,N}, labels) where
     )
     data = ReactantCore.materialize_traced_array(reshape(data, 1, size(data)...))
     indices = UInt32.(@opcall(findfirst(data .== labels_expanded; dimension=1)))
-    return OneHotArray{TracedRNumber{UInt32},N,N + 1,typeof(indices)}(
+    return OneHotArray{Reactant.TracedRInteger{UInt32},N,N + 1,typeof(indices)}(
         indices, length(labels)
     )
 end
@@ -73,7 +77,7 @@ function OneHotArrays.onehotbatch(
         TracedRNumber{UInt32} ∘ Base.Fix2(+, 1 - first(labels)),
         ReactantCore.materialize_traced_array(data),
     )
-    return OneHotArray{TracedRNumber{UInt32},N,N + 1,typeof(indices)}(
+    return OneHotArray{Reactant.TracedRInteger{UInt32},N,N + 1,typeof(indices)}(
         indices, length(labels)
     )
 end
diff --git a/ext/ReactantSparseArraysExt/ReadOnly.jl b/ext/ReactantSparseArraysExt/ReadOnly.jl
index 6c8630300c..3903a405fb 100644
--- a/ext/ReactantSparseArraysExt/ReadOnly.jl
+++ b/ext/ReactantSparseArraysExt/ReadOnly.jl
@@ -49,7 +49,7 @@ function Base.getindex(
 end
 
 function Base.getindex(
-    x::SparseArrays.ReadOnly{Reactant.TracedRNumber{T},1,V},
+    x::SparseArrays.ReadOnly{<:Reactant.TracedRNumber{T},1,V},
     idx::Union{Int64,Reactant.TracedRNumber{Int64}},
 ) where {T,V<:AbstractArray{<:Reactant.TracedRNumber{T},1}}
     return getindex(parent(x), idx)
diff --git a/src/stdlibs/LinearAlgebra.jl b/src/stdlibs/LinearAlgebra.jl
index 1c4b343763..2c7bbb2ca9 100644
--- a/src/stdlibs/LinearAlgebra.jl
+++ b/src/stdlibs/LinearAlgebra.jl
@@ -118,7 +118,7 @@ for (AT, comp) in ((:LowerTriangular, "GE"), (:UpperTriangular, "LE"))
             px = materialize_traced_array(parent(x))
             row_idxs = @opcall iota(Int, [m, n]; iota_dimension=1)
             col_idxs = @opcall iota(Int, [m, n]; iota_dimension=2)
-            indicator = @opcall compare(row_idxs, col_idxs; comparison_direction=$(comp))
+            indicator = @opcall compare(row_idxs, col_idxs; comparison_direction=($(comp)))
             return @opcall select(indicator, px, zero(px))
         end
 
@@ -214,9 +214,11 @@ for (AT, dcomp, ocomp) in (
         m, n = size(x)
         row_idxs = @opcall iota(Int, [m, n]; iota_dimension=1)
         col_idxs = @opcall iota(Int, [m, n]; iota_dimension=2)
-        data_indicator = @opcall compare(row_idxs, col_idxs; comparison_direction=$(dcomp))
+        data_indicator = @opcall compare(
+            row_idxs, col_idxs; comparison_direction=($(dcomp))
+        )
         original_indicator = @opcall compare(
-            row_idxs, col_idxs; comparison_direction=$(ocomp)
+            row_idxs, col_idxs; comparison_direction=($(ocomp))
         )
         res = @opcall add(
             @opcall(select(data_indicator, tdata, z)),
@@ -784,9 +786,9 @@ for (wT, lower, ud) in (
             parent(A),
             rhs;
             left_side=false,
-            lower=$(lower),
-            transpose_a='N',
-            unit_diagonal=$(ud),
+            lower=($(lower)),
+            transpose_a=('N'),
+            unit_diagonal=($(ud)),
         )
     end
 end

From c9ecf7ba030caa73e0ad5d0d5851d475b4dc308d Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Tue, 7 Jul 2026 20:41:10 +0200
Subject: [PATCH 09/30] Resolve method ambiguities against Base's specialized
 number methods

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.
---
 src/ConcreteRArray.jl | 134 ++++++++++++++++++++++------
 src/Indexing.jl       |  40 +++++----
 src/TracedRArray.jl   |  11 ++-
 src/TracedRNumber.jl  | 202 +++++++++++++++++++++++++++++++++++-------
 src/TracedRational.jl |  80 +++++++++++------
 5 files changed, 361 insertions(+), 106 deletions(-)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 4ca5afe420..0ba52e86fe 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -243,8 +243,18 @@ for jlop in (
         for T2 in CONCRETE_NUMBER_KINDS
             @eval $(jlop)(x::$(T1), y::$(T2)) = $(jlop)(to_number(x), to_number(y))
         end
-        for X in
-            (:Rational, :BigInt, :BigFloat, :AbstractIrrational, :AbstractFloat, :Complex)
+        for X in (
+            :Real,
+            :Integer,
+            :Bool,
+            :Rational,
+            :BigInt,
+            :BigFloat,
+            :AbstractIrrational,
+            :AbstractFloat,
+            :Complex,
+            :(Complex{Bool}),
+        )
             @eval begin
                 $(jlop)(x::$(T1), y::$(X)) = $(jlop)(to_number(x), y)
                 $(jlop)(x::$(X), y::$(T1)) = $(jlop)(x, to_number(y))
@@ -269,33 +279,95 @@ for jlop in (
     end
 end
 
-for T in (AbstractConcreteNumber, AbstractConcreteArray{<:Number,0})
-    @eval Base.div(x::$(T), y::$(T), r::RoundingMode=RoundToZero) =
-        div(to_number(x), to_number(y), r)
-    @eval Base.div(x::$(T), y::Number, r::RoundingMode=RoundToZero) =
-        div(to_number(x), y, r)
-    @eval Base.div(x::Number, y::$(T), r::RoundingMode=RoundToZero) =
-        div(x, to_number(y), r)
+# Base specializes `div` on single rounding modes and on `Rational` arguments;
+# the corresponding methods below only disambiguate against those.
+const BASE_SPECIFIC_ROUNDING_MODES = (
+    RoundingMode{:FromZero},
+    RoundingMode{:Nearest},
+    RoundingMode{:NearestTiesAway},
+    RoundingMode{:NearestTiesUp},
+    RoundingMode{:Up},
+    RoundingMode{:Down},
+    # Base also groups the nearest modes into a single method
+    Union{
+        RoundingMode{:Nearest},RoundingMode{:NearestTiesAway},RoundingMode{:NearestTiesUp}
+    },
+)
+
+for T1 in CONCRETE_NUMBER_KINDS
+    @eval begin
+        Base.div(x::$(T1), y::Number, r::RoundingMode=RoundToZero) =
+            div(to_number(x), y, r)
+        Base.div(x::Number, y::$(T1), r::RoundingMode=RoundToZero) =
+            div(x, to_number(y), r)
+        Base.div(x::$(T1), y::Real, r::RoundingMode=RoundToZero) = div(to_number(x), y, r)
+        Base.div(x::Real, y::$(T1), r::RoundingMode=RoundToZero) = div(x, to_number(y), r)
+        Base.div(x::$(T1), y::Rational, r::RoundingMode=RoundToZero) =
+            div(to_number(x), y, r)
+        Base.div(x::Rational, y::$(T1), r::RoundingMode=RoundToZero) =
+            div(x, to_number(y), r)
+        Base.div(x::$(T1), y::TracedRReal, r::RoundingMode=RoundToZero) =
+            div(to_number(x), y, r)
+        Base.div(x::TracedRReal, y::$(T1), r::RoundingMode=RoundToZero) =
+            div(x, to_number(y), r)
+    end
+    for RM in BASE_SPECIFIC_ROUNDING_MODES
+        @eval begin
+            Base.div(x::$(T1), y::Integer, r::$RM) = div(to_number(x), y, r)
+            Base.div(x::Integer, y::$(T1), r::$RM) = div(x, to_number(y), r)
+            Base.div(x::T, y::T, r::$RM) where {T<:$(T1)} =
+                div(to_number(x), to_number(y), r)
+        end
+    end
+    for T2 in CONCRETE_NUMBER_KINDS
+        @eval Base.div(x::$(T1), y::$(T2), r::RoundingMode=RoundToZero) =
+            div(to_number(x), to_number(y), r)
+        for RM in BASE_SPECIFIC_ROUNDING_MODES
+            @eval Base.div(x::$(T1), y::$(T2), r::$RM) =
+                div(to_number(x), to_number(y), r)
+        end
+    end
+    for T2 in TRACED_NUMBER_KINDS
+        @eval begin
+            Base.div(x::$(T1), y::$(T2), r::RoundingMode=RoundToZero) =
+                div(to_number(x), y, r)
+            Base.div(x::$(T2), y::$(T1), r::RoundingMode=RoundToZero) =
+                div(x, to_number(y), r)
+        end
+        for RM in BASE_SPECIFIC_ROUNDING_MODES
+            @eval begin
+                Base.div(x::$(T1), y::$(T2), r::$RM) = div(to_number(x), y, r)
+                Base.div(x::$(T2), y::$(T1), r::$RM) = div(x, to_number(y), r)
+            end
+        end
+    end
+end
 
-    @eval Base.div(x::$(T), y::TracedRNumber, r::RoundingMode=RoundToZero) =
-        div(to_number(x), y, r)
-    @eval Base.div(x::TracedRNumber, y::$(T), r::RoundingMode=RoundToZero) =
-        div(x, to_number(y), r)
-    @eval Base.div(x::TracedRNumber{S}, y::$(T), r::RoundingMode) where {S<:Integer} =
-        div(x, to_number(y), r)
-    @eval Base.div(x::TracedRNumber{S}, y::$(T), r::RoundingMode) where {S<:AbstractFloat} =
-        div(x, to_number(y), r)
-    @eval Base.div(x::$(T), y::TracedRNumber{S}, r::RoundingMode) where {S<:Integer} =
-        div(to_number(x), y, r)
-    @eval Base.div(x::$(T), y::TracedRNumber{S}, r::RoundingMode) where {S<:AbstractFloat} =
-        div(to_number(x), y, r)
+let T = AbstractConcreteArray{<:Number,0}
+    @eval begin
+        Base.div(x::$(T), y::$(T), r::RoundingMode=RoundToZero) =
+            div(to_number(x), to_number(y), r)
+        Base.div(x::$(T), y::Number, r::RoundingMode=RoundToZero) = div(to_number(x), y, r)
+        Base.div(x::Number, y::$(T), r::RoundingMode=RoundToZero) = div(x, to_number(y), r)
+        Base.div(x::$(T), y::TracedRNumber, r::RoundingMode=RoundToZero) =
+            div(to_number(x), y, r)
+        Base.div(x::TracedRNumber, y::$(T), r::RoundingMode=RoundToZero) =
+            div(x, to_number(y), r)
+    end
 end
 
-# `Rational` exponents are covered by the disambiguation grid above.
+# `Rational` and `Integer` exponents are covered by the disambiguation grid
+# above; these disambiguate against Base's `^(::SomeFloat, ::Integer)` and
+# `^(::Complex{...}, ::Integer)` specializations.
+for B in (Float16, Float32, Union{Float16,Float32}, Float64, BFloat16)
+    @eval Base.:^(x::$(B), y::AbstractConcreteInteger) = ^(x, to_number(y))
+end
+for B in (Complex{<:AbstractFloat}, Complex{<:Integer}, Complex{<:Rational})
+    @eval Base.:^(x::$(B), y::AbstractConcreteInteger) = ^(x, to_number(y))
+end
 for CT in CONCRETE_NUMBER_KINDS
-    @eval Base.:^(x::$(CT), y::Integer) = ^(to_number(x), y)
+    @eval Base.:^(::Irrational{:ℯ}, x::$(CT)) = exp(to_number(x))
 end
-Base.:^(::Irrational{:ℯ}, x::AbstractConcreteNumber) = exp(to_number(x))
 
 for jlop in (:(Base.isnan), :(Base.isfinite)),
     T in (CONCRETE_NUMBER_KINDS..., AbstractConcreteArray{<:Any,0})
@@ -303,13 +375,17 @@ for jlop in (:(Base.isnan), :(Base.isfinite)),
     @eval $(jlop)(x::$(T)) = $(jlop)(to_number(x))
 end
 
-for (T1, T2) in (
-    (AbstractConcreteNumber, AbstractConcreteNumber),
-    (AbstractConcreteNumber, Number),
-    (Number, AbstractConcreteNumber),
+isapprox_pairs = Any[
     (AbstractConcreteArray{<:Any,0}, Number),
     (Number, AbstractConcreteArray{<:Any,0}),
-)
+]
+for T1 in CONCRETE_NUMBER_KINDS
+    push!(isapprox_pairs, (T1, Number), (Number, T1), (T1, Integer), (Integer, T1))
+    for T2 in CONCRETE_NUMBER_KINDS
+        push!(isapprox_pairs, (T1, T2))
+    end
+end
+for (T1, T2) in isapprox_pairs
     @eval begin
         function Base.isapprox(x::$(T1), y::$(T2); kwargs...)
             return Base.isapprox(to_number(x), to_number(y); kwargs...)
diff --git a/src/Indexing.jl b/src/Indexing.jl
index 3b2156216b..e2b6a7a4d7 100644
--- a/src/Indexing.jl
+++ b/src/Indexing.jl
@@ -1,6 +1,12 @@
 module TracedIndexing
 
-using ..Reactant: Reactant, TracedRArray, TracedRNumber, TracedStepRangeLen, TracedUnitRange
+using ..Reactant:
+    Reactant,
+    TracedRArray,
+    TracedRNumber,
+    TracedRInteger,
+    TracedStepRangeLen,
+    TracedUnitRange
 using ..Reactant: AnyTracedRArray, ancestor, unwrapped_eltype
 using ..Ops: @opcall
 using ..TracedUtils: TracedUtils
@@ -83,10 +89,11 @@ function Base.getindex(
 end
 
 for (aT, iT) in (
-    (AbstractRange, TracedRNumber{<:Integer}),
+    (AbstractRange, TracedRInteger),
     (AbstractRange{<:TracedRNumber}, Int),
-    (AbstractRange{<:TracedRNumber}, TracedRNumber{<:Integer}),
-    (AbstractRange{<:TracedRNumber}, TracedRNumber{Int}),
+    (AbstractRange{<:TracedRNumber}, TracedRInteger),
+    (AbstractRange{<:TracedRNumber}, Union{Int,TracedRInteger{Int}}),
+    (Base.IdentityUnitRange, TracedRInteger),
     # 1.10 specific ambiguity fixes
     (UnitRange{<:TracedRNumber}, Int),
 )
@@ -99,10 +106,9 @@ for (aT, iT) in (
 end
 
 for (aT, iT) in (
-    (Base.OneTo, TracedRNumber{<:Integer}),
+    (Base.OneTo, TracedRInteger),
     (Base.OneTo{<:TracedRNumber}, Int),
-    (Base.OneTo{<:TracedRNumber}, TracedRNumber{<:Integer}),
-    (Base.OneTo{<:TracedRNumber}, TracedRNumber{Int}),
+    (Base.OneTo{<:TracedRNumber}, TracedRInteger),
 )
     @eval function Base.getindex(a::$aT, index::$iT)
         return convert(TracedRNumber{Reactant.unwrapped_eltype(a)}, index)
@@ -110,10 +116,9 @@ for (aT, iT) in (
 end
 
 for (aT, iT) in (
-    (StepRangeLen, TracedRNumber{<:Integer}),
+    (StepRangeLen, TracedRInteger),
     (StepRangeLen{<:TracedRNumber}, Int),
-    (StepRangeLen{<:TracedRNumber}, TracedRNumber{<:Integer}),
-    (StepRangeLen{<:TracedRNumber}, TracedRNumber{Int}),
+    (StepRangeLen{<:TracedRNumber}, TracedRInteger),
 )
     @eval function Base.getindex(r::$aT, index::$iT)
         # FIXME: this crashes for some reason
@@ -126,10 +131,9 @@ for (aT, iT) in (
 end
 
 for (aT, iT) in (
-    (LinRange, TracedRNumber{<:Integer}),
+    (LinRange, TracedRInteger),
     (LinRange{<:TracedRNumber}, Int),
-    (LinRange{<:TracedRNumber}, TracedRNumber{<:Integer}),
-    (LinRange{<:TracedRNumber}, TracedRNumber{Int}),
+    (LinRange{<:TracedRNumber}, TracedRInteger),
 )
     @eval function Base.getindex(r::$aT, index::$iT)
         return convert(
@@ -290,10 +294,12 @@ function Base.unsafe_getindex(r::TracedStepRangeLen, i::TracedRNumber{Int})
 end
 
 function Base.unsafe_getindex(
-    r::Union{
-        Base.StepRangeLen{T,<:TwicePrecision,<:TwicePrecision},
-        TracedStepRangeLen{T,<:TwicePrecision,<:TwicePrecision,<:TwicePrecision},
-    },
+    r::Base.StepRangeLen{T,<:TwicePrecision,<:TwicePrecision}, i::TracedRInteger{Int}
+) where {T}
+    return overloaded_unsafe_getindex(r, i)
+end
+function Base.unsafe_getindex(
+    r::TracedStepRangeLen{T,<:TwicePrecision,<:TwicePrecision,<:TwicePrecision},
     i::TracedRNumber{Int},
 ) where {T}
     return overloaded_unsafe_getindex(r, i)
diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
index 77a36e2d70..ad12a279e0 100644
--- a/src/TracedRArray.jl
+++ b/src/TracedRArray.jl
@@ -3,7 +3,13 @@ module TracedRArrayOverrides
 using Base: Broadcast
 using Base.Broadcast: Broadcasted, AbstractArrayStyle, instantiate
 
-using ..Reactant: Reactant, TracedRArray, TracedRNumber, AnyTracedRArray, AnyTracedRVector
+using ..Reactant:
+    Reactant,
+    TracedRArray,
+    TracedRNumber,
+    TracedRInteger,
+    AnyTracedRArray,
+    AnyTracedRVector
 using ..Reactant: MLIR, unwrapped_eltype
 using ..Ops: @opcall
 using ..TracedUtils: TracedUtils, get_mlir_data, set_mlir_data!, materialize_traced_array
@@ -254,6 +260,9 @@ end
 function Base.fill!(A::Array{T,N}, x::TracedRNumber{T2}) where {T,N,T2}
     return throw(MethodError(fill!, (A, x)))
 end
+function Base.fill!(A::Union{Array{Int8},Array{UInt8}}, x::TracedRInteger)
+    return throw(MethodError(fill!, (A, x)))
+end
 
 struct AbstractReactantArrayStyle{N} <: AbstractArrayStyle{N} end
 
diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index 9d2ed82601..855116d0f6 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -111,8 +111,19 @@ function Base.promote_rule(::Type{<:TracedRNumber{T}}, ::Type{Bool}) where {T}
     return traced_number_type(T)
 end
 
-# Disambiguate against Base's irrational-vs-Real rule; real-only since
-# `TracedRReal <: Real` while the full union is not.
+# Disambiguate against Base's irrational rules on `Number` and `Real`.
+function Base.promote_rule(
+    T::Type{<:AbstractIrrational}, ::Type{<:TracedRNumber{S}}
+) where {S}
+    return traced_number_type(Base.promote_type(T, S))
+end
+
+function Base.promote_rule(
+    ::Type{<:TracedRNumber{S}}, T::Type{<:AbstractIrrational}
+) where {S}
+    return traced_number_type(Base.promote_type(T, S))
+end
+
 function Base.promote_rule(
     T::Type{<:AbstractIrrational}, ::Type{<:TracedRReal{S}}
 ) where {S}
@@ -125,6 +136,36 @@ function Base.promote_rule(
     return traced_number_type(Base.promote_type(T, S))
 end
 
+# Disambiguate against the float8 promotion rules in PrimitiveTypes.jl, which
+# apply to `Type{<:Integer}` and hence to traced integers.
+for S in Base.uniontypes(Reactant.ReactantFloat8)
+    @eval begin
+        function Base.promote_rule(::Type{$S}, ::Type{<:TracedRInteger{T}}) where {T}
+            return traced_number_type(Base.promote_type(T, $S))
+        end
+        function Base.promote_rule(::Type{<:TracedRInteger{T}}, ::Type{$S}) where {T}
+            return traced_number_type(Base.promote_type(T, $S))
+        end
+    end
+end
+
+# Disambiguate against Base's `Rational` and `Complex` promotion rules.
+function Base.promote_rule(
+    ::Type{Rational{T}}, ::Type{<:TracedRInteger{S}}
+) where {T<:Integer,S}
+    return Reactant.TracedRational{traced_number_type(Base.promote_type(T, S))}
+end
+function Base.promote_rule(
+    ::Type{Rational{T}}, ::Type{<:TracedRFloat{S}}
+) where {T<:Integer,S}
+    return traced_number_type(Base.promote_type(Rational{T}, S))
+end
+function Base.promote_rule(
+    ::Type{Complex{T}}, ::Type{<:TracedRReal{S}}
+) where {T<:Real,S}
+    return traced_number_type(Base.promote_type(Complex{T}, S))
+end
+
 function Base.promote_rule(::Type{Nothing}, ::Type{<:TracedRNumber{S}}) where {S}
     return Union{Nothing,TracedRNumber{S}}
 end
@@ -206,7 +247,18 @@ function TracedRNumber{T}(x::Number) where {T}
 end
 
 # Covers the concrete traced types, e.g. via `convert(TracedRFloat{Float64}, x)`.
+# The extra methods disambiguate against Base's constructors on `Real`,
+# `AbstractFloat`, and `Integer`.
 (::Type{RT})(x::Number) where {RT<:TracedRNumber} = Reactant.promote_to(RT, x)
+for (RT, XT) in (
+    (:TracedRReal, :Complex),
+    (:TracedRInteger, :Complex),
+    (:TracedRFloat, :Complex),
+    (:TracedRInteger, :Rational),
+    (:TracedRFloat, :Rational),
+)
+    @eval (::Type{RT})(x::$XT) where {RT<:$RT} = Reactant.promote_to(RT, x)
+end
 
 for T in Base.uniontypes(Reactant.ReactantFloat8)
     @eval TracedRNumber{T}(x::$T) where {T} = Reactant.promote_to(TracedRNumber{T}, x)
@@ -312,6 +364,21 @@ for op in (:mod, :mod1, :rem)
     end
 end
 
+for op in (:mod, :rem)
+    @eval begin
+        function Base.$op(
+            @nospecialize(lhs::TracedRInteger{T}), @nospecialize(rhs::Rational)
+        ) where {T}
+            return $(op)(lhs, Reactant.promote_to(TracedRNumber{T}, rhs))
+        end
+        function Base.$op(
+            @nospecialize(lhs::Rational), @nospecialize(rhs::TracedRInteger{T})
+        ) where {T}
+            return $(op)(Reactant.promote_to(TracedRNumber{T}, lhs), rhs)
+        end
+    end
+end
+
 Base.flipsign(x::TracedRReal, y::TracedRReal) = ifelse(y < 0, -x, x)
 
 function Base.div(
@@ -405,30 +472,30 @@ function Base.div(
 end
 
 function Base.div(
-    @nospecialize(lhs::TracedRNumber{T}),
-    @nospecialize(rhs::TracedRNumber{T}),
+    @nospecialize(lhs::TracedRInteger{T}),
+    @nospecialize(rhs::TracedRInteger{T}),
     ::typeof(RoundToZero),
 ) where {T<:Integer}
     return @opcall divide(lhs, rhs)
 end
 function Base.div(
-    @nospecialize(lhs::TracedRNumber{T}),
-    @nospecialize(rhs::TracedRNumber{T}),
+    @nospecialize(lhs::TracedRInteger{T}),
+    @nospecialize(rhs::TracedRInteger{T}),
     ::typeof(RoundDown),
 ) where {T<:Integer}
     return @opcall divide(lhs, rhs)
 end
 function Base.div(
-    @nospecialize(lhs::TracedRNumber{T}),
-    @nospecialize(rhs::TracedRNumber{T}),
+    @nospecialize(lhs::TracedRInteger{T}),
+    @nospecialize(rhs::TracedRInteger{T}),
     ::typeof(RoundUp),
 ) where {T<:Integer}
     q = div(lhs, rhs)  # truncation (RoundToZero)
     return q + (!iszero(rem(lhs, rhs)) & (signbit(lhs) == signbit(rhs)))
 end
 function Base.div(
-    @nospecialize(lhs::TracedRNumber{T}),
-    @nospecialize(rhs::TracedRNumber{T}),
+    @nospecialize(lhs::TracedRInteger{T}),
+    @nospecialize(rhs::TracedRInteger{T}),
     ::typeof(RoundFromZero),
 ) where {T<:Integer}
     return ifelse(
@@ -436,25 +503,81 @@ function Base.div(
     )
 end
 function Base.div(
-    @nospecialize(lhs::TracedRNumber{T}),
-    @nospecialize(rhs::TracedRNumber{T}),
+    @nospecialize(lhs::TracedRInteger{T}),
+    @nospecialize(rhs::TracedRInteger{T}),
     r::Union{typeof(RoundNearest),typeof(RoundNearestTiesAway),typeof(RoundNearestTiesUp)},
 ) where {T<:Integer}
     return divrem(lhs, rhs, r)[1]
 end
 
 function Base.div(
-    @nospecialize(x::TracedRNumber{T}), @nospecialize(y::TracedRNumber{T}), r::RoundingMode
-) where {T<:AbstractFloat}
+    @nospecialize(x::TracedRFloat{T}), @nospecialize(y::TracedRFloat{T}), r::RoundingMode
+) where {T}
     return round(div(x, y), r)
 end
 
-Base.div(@nospecialize(lhs::TracedRNumber{<:Integer}), ::Missing, ::RoundingMode) = missing
-Base.div(::Missing, @nospecialize(rhs::TracedRNumber{<:Integer}), ::RoundingMode) = missing
+Base.div(@nospecialize(lhs::TracedRInteger), ::Missing, ::RoundingMode) = missing
+Base.div(::Missing, @nospecialize(rhs::TracedRInteger), ::RoundingMode) = missing
+
+# Base has `div` methods specialized on single rounding modes (and on
+# `Rational` arguments); the following methods only disambiguate against them.
+for RM in (
+    RoundingMode{:FromZero},
+    RoundingMode{:Nearest},
+    RoundingMode{:NearestTiesAway},
+    RoundingMode{:NearestTiesUp},
+    RoundingMode{:Up},
+    RoundingMode{:Down},
+    # Base also groups the nearest modes into a single method
+    Union{
+        RoundingMode{:Nearest},RoundingMode{:NearestTiesAway},RoundingMode{:NearestTiesUp}
+    },
+)
+    @eval begin
+        function Base.div(
+            @nospecialize(lhs::TracedRInteger{T1}),
+            @nospecialize(rhs::TracedRInteger{T2}),
+            r::$RM,
+        ) where {T1,T2}
+            T = promote_type(T1, T2)
+            return div(
+                Reactant.promote_to(TracedRNumber{T}, lhs),
+                Reactant.promote_to(TracedRNumber{T}, rhs),
+                r,
+            )
+        end
+        function Base.div(
+            @nospecialize(lhs::TracedRInteger{T}), rhs::Integer, r::$RM
+        ) where {T}
+            return div(lhs, Reactant.promote_to(TracedRNumber{T}, rhs), r)
+        end
+        function Base.div(
+            lhs::Integer, @nospecialize(rhs::TracedRInteger{T}), r::$RM
+        ) where {T}
+            return div(Reactant.promote_to(TracedRNumber{T}, lhs), rhs, r)
+        end
+    end
+end
+function Base.div(
+    @nospecialize(lhs::TracedRInteger{T}), rhs::Rational, r::RoundingMode
+) where {T}
+    return div(lhs, Reactant.promote_to(TracedRNumber{T}, rhs), r)
+end
+function Base.div(
+    lhs::Rational, @nospecialize(rhs::TracedRInteger{T}), r::RoundingMode
+) where {T}
+    return div(Reactant.promote_to(TracedRNumber{T}, lhs), rhs, r)
+end
+function Base.div(@nospecialize(lhs::TracedRInteger{T}), rhs::Rational) where {T}
+    return div(lhs, Reactant.promote_to(TracedRNumber{T}, rhs))
+end
+function Base.div(lhs::Rational, @nospecialize(rhs::TracedRInteger{T})) where {T}
+    return div(Reactant.promote_to(TracedRNumber{T}, lhs), rhs)
+end
 
 function Base.divrem(
-    @nospecialize(a::TracedRNumber{T}),
-    @nospecialize(b::TracedRNumber{T}),
+    @nospecialize(a::TracedRInteger{T}),
+    @nospecialize(b::TracedRInteger{T}),
     r::Union{typeof(RoundUp),typeof(RoundDown),typeof(RoundToZero)},
 ) where {T<:Integer}
     if r === RoundToZero
@@ -470,10 +593,10 @@ function Base.divrem(
 end
 
 function Base.divrem(
-    @nospecialize(x::TracedRNumber{T}),
-    @nospecialize(y::TracedRNumber{T}),
+    @nospecialize(x::TracedRInteger{T}),
+    @nospecialize(y::TracedRInteger{T}),
     ::typeof(RoundNearest),
-) where {T<:Integer}
+) where {T}
     (q, r) = divrem(x, y)
     threshold = isodd(y) | iseven(q)
     half_y = y ÷ 2
@@ -492,10 +615,10 @@ function Base.divrem(
 end
 
 function Base.divrem(
-    @nospecialize(x::TracedRNumber{T}),
-    @nospecialize(y::TracedRNumber{T}),
+    @nospecialize(x::TracedRInteger{T}),
+    @nospecialize(y::TracedRInteger{T}),
     ::typeof(RoundNearestTiesAway),
-) where {T<:Integer}
+) where {T}
     (q, r) = divrem(x, y)
     threshold = isodd(y)
     half_y = y ÷ 2
@@ -514,10 +637,10 @@ function Base.divrem(
 end
 
 function Base.divrem(
-    @nospecialize(x::TracedRNumber{T}),
-    @nospecialize(y::TracedRNumber{T}),
+    @nospecialize(x::TracedRInteger{T}),
+    @nospecialize(y::TracedRInteger{T}),
     ::typeof(RoundNearestTiesUp),
-) where {T<:Integer}
+) where {T}
     (q, r) = divrem(x, y)
     half_y = y ÷ 2
     # x >= 0, y >= 0
@@ -535,10 +658,10 @@ function Base.divrem(
 end
 
 function Base.divrem(
-    @nospecialize(x::TracedRNumber{T}),
-    @nospecialize(y::TracedRNumber{T}),
+    @nospecialize(x::TracedRInteger{T}),
+    @nospecialize(y::TracedRInteger{T}),
     ::typeof(RoundFromZero),
-) where {T<:Integer}
+) where {T}
     q_up, r_up = divrem(x, y, RoundUp)
     q_down, r_down = divrem(x, y, RoundDown)
     return ifelse(signbit(x) == signbit(y), (q_up, r_up), (q_down, r_down))
@@ -689,6 +812,12 @@ function Base.:+(
     return Base.TwicePrecision(Base.canonicalize2(r, s)...)
 end
 
+function Base.:*(
+    x::TwicePrecision{<:Union{Float16,Float32,Float64}}, v::TracedRInteger
+)
+    return invoke(*, Tuple{TwicePrecision,TracedRNumber}, x, v)
+end
+
 function Base.:*(x::TwicePrecision, v::TracedRNumber)
     @trace result = if v == 0
         TwicePrecision(x.hi * v, x.lo * v)
@@ -1059,6 +1188,13 @@ function Base.copysign(x::S, y::TracedRReal) where {S<:Real}
     return copysign(Reactant.promote_to(TracedRNumber{unwrapped_eltype(S)}, x), y)
 end
 
+# Base specializes `copysign` on these first-argument types.
+for S in (:Float16, :Float32, :Float64, :Signed, :Rational)
+    @eval function Base.copysign(x::$S, y::TracedRReal)
+        return copysign(Reactant.promote_to(TracedRNumber{unwrapped_eltype(x)}, x), y)
+    end
+end
+
 function Base.zeros(::Type{<:TracedRNumber{T}}, dims::Dims{N}) where {T,N}
     return @opcall fill(zero(T), dims)
 end
@@ -1093,6 +1229,12 @@ function Base.checkindex(::Type{Bool}, ::AbstractUnitRange, ::TracedRNumber)
     return true
 end
 
+function Base.checkindex(::Type{Bool}, ::Base.IdentityUnitRange, ::TracedRReal)
+    @warn "Currently we don't perform bounds checking for TracedRNumber. This will be \
+           fixed in a future version of Reactant." maxlog = 1
+    return true
+end
+
 # rem2pi: fallback to rem for now.
 # TODO(#2259): we should replace with the more numerically stable version.
 # https://github.com/JuliaLang/julia/blob/4d04bb6b3b1b879f4dbb918d194c5c939a1e7f3c/base/special/rem2pi.jl#L133
diff --git a/src/TracedRational.jl b/src/TracedRational.jl
index 88b80fe2fb..a540367a28 100644
--- a/src/TracedRational.jl
+++ b/src/TracedRational.jl
@@ -1,6 +1,12 @@
 module TracedRationalOverrides
 
-using ..Reactant: Reactant, AbstractConcreteNumber, TracedRNumber, TracedRInteger, RInteger
+using ..Reactant:
+    Reactant,
+    AbstractConcreteNumber,
+    AbstractConcreteInteger,
+    TracedRNumber,
+    TracedRInteger,
+    RInteger
 import ..Reactant: TracedRational
 using ReactantCore: ReactantCore
 
@@ -118,18 +124,19 @@ Base.:-(x::TracedRational) = TracedRational(-x.num, x.den)
 # TODO: check for overflow/underflow for integer types
 
 for op in (:+, :-, :rem, :mod)
-    @eval begin
-        function Base.$(op)(x::TracedRational, y::TracedRational)
-            xd, yd = Base.divgcd(promote(x.den, y.den)...)
-            return TracedRational($(op)(x.num * yd, y.num * xd), x.den * yd)
-        end
-
-        function Base.$(op)(x::TracedRational, y::Integer)
-            return TracedRational($(op)(x.num, x.den * y), x.den)
-        end
-
-        function Base.$(op)(x::Integer, y::TracedRational)
-            return TracedRational($(op)(x * y.den, y.num), y.den)
+    @eval function Base.$(op)(x::TracedRational, y::TracedRational)
+        xd, yd = Base.divgcd(promote(x.den, y.den)...)
+        return TracedRational($(op)(x.num * yd, y.num * xd), x.den * yd)
+    end
+    for IT in (:Integer, :AbstractConcreteInteger)
+        @eval begin
+            function Base.$(op)(x::TracedRational, y::$IT)
+                return TracedRational($(op)(x.num, x.den * y), x.den)
+            end
+
+            function Base.$(op)(x::$IT, y::TracedRational)
+                return TracedRational($(op)(x * y.den, y.num), y.den)
+            end
         end
     end
 end
@@ -140,14 +147,18 @@ function Base.:*(x::TracedRational, y::TracedRational)
     return TracedRational(xn * yn, xd * yd)
 end
 
-function Base.:*(x::TracedRational, y::Integer)
-    xd, yn = Base.divgcd(promote(x.den, y)...)
-    return TracedRational(x.num * yn, xd)
-end
+for IT in (:Integer, :AbstractConcreteInteger)
+    @eval begin
+        function Base.:*(x::TracedRational, y::$IT)
+            xd, yn = Base.divgcd(promote(x.den, y)...)
+            return TracedRational(x.num * yn, xd)
+        end
 
-function Base.:*(x::Integer, y::TracedRational)
-    xn, yd = Base.divgcd(promote(x, y.den)...)
-    return TracedRational(xn * y.num, yd)
+        function Base.:*(x::$IT, y::TracedRational)
+            xn, yd = Base.divgcd(promote(x, y.den)...)
+            return TracedRational(xn * y.num, yd)
+        end
+    end
 end
 
 Base.inv(x::TracedRational) = TracedRational(x.den, x.num)
@@ -165,6 +176,9 @@ Base.convert(::Type{TracedRational{T}}, x::TracedRational{T}) where {T} = x
 Base.convert(::Type{TracedRational{T}}, x::TracedRational) where {T} = TracedRational{T}(x)
 Base.convert(::Type{TracedRational{T}}, x::Rational) where {T} = TracedRational{T}(x)
 Base.convert(::Type{TracedRational{T}}, x::Integer) where {T} = TracedRational{T}(x)
+function Base.convert(::Type{TracedRational{T}}, x::AbstractConcreteInteger) where {T}
+    return TracedRational{T}(x)
+end
 
 # Other utility functions
 Base.zero(::Type{TracedRational{T}}) where {T} = TracedRational(zero(T), one(T))
@@ -189,14 +203,18 @@ end
 Base.://(num::TracedRational, den::Rational) = num//TracedRational(den)
 Base.://(num::Rational, den::TracedRational) = TracedRational(num)//den
 
-function Base.://(x::TracedRational, y::Integer)
-    xn, yn = Base.divgcd(promote(x.num, y)...)
-    return checked_den(xn, x.den * yn)
-end
+for IT in (:Integer, :AbstractConcreteInteger)
+    @eval begin
+        function Base.://(x::TracedRational, y::$IT)
+            xn, yn = Base.divgcd(promote(x.num, y)...)
+            return checked_den(xn, x.den * yn)
+        end
 
-function Base.://(x::Integer, y::TracedRational)
-    xn, yn = Base.divgcd(promote(x, y.num)...)
-    return checked_den(xn * y.den, yn)
+        function Base.://(x::$IT, y::TracedRational)
+            xn, yn = Base.divgcd(promote(x, y.num)...)
+            return checked_den(xn * y.den, yn)
+        end
+    end
 end
 
 function Base.://(x::TracedRational, y::TracedRational)
@@ -205,8 +223,12 @@ function Base.://(x::TracedRational, y::TracedRational)
     return checked_den(xn * yd, xd * yn)
 end
 
-Base.:/(x::TracedRational, y::Integer) = x//y
-Base.:/(x::Integer, y::TracedRational) = x//y
+for IT in (:Integer, :AbstractConcreteInteger)
+    @eval begin
+        Base.:/(x::TracedRational, y::$IT) = x//y
+        Base.:/(x::$IT, y::TracedRational) = x//y
+    end
+end
 Base.:/(x::TracedRational, y::TracedRational) = x//y
 
 end

From db04e94f7ba5aed584afa1edda2f826a5db6e1af Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Tue, 7 Jul 2026 21:54:58 +0200
Subject: [PATCH 10/30] Fix test-suite fallout from the number-kind split

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.
---
 src/ConcreteRArray.jl |  4 ++++
 src/Reactant.jl       |  5 +++++
 src/TracedRNumber.jl  | 21 ++++++++++++++++++---
 src/TracedRational.jl | 10 ++++++++++
 src/TracedUtils.jl    |  4 ++++
 src/Tracing.jl        | 10 +++++++++-
 src/Types.jl          | 19 +++++++++++++++++++
 test/core/tracing.jl  | 14 +++++++++-----
 8 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 0ba52e86fe..97483e4c11 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -228,6 +228,10 @@ const TRACED_NUMBER_KINDS = (TracedRInteger, TracedRFloat, TracedRComplex)
 
 for jlop in (
     :(Base.isless),
+    :(Base.:<),
+    :(Base.:<=),
+    :(Base.:>),
+    :(Base.:>=),
     :(Base.:+),
     :(Base.:-),
     :(Base.:*),
diff --git a/src/Reactant.jl b/src/Reactant.jl
index f02978433b..2ba6d0f790 100644
--- a/src/Reactant.jl
+++ b/src/Reactant.jl
@@ -159,6 +159,8 @@ aos_to_soa(x::AbstractArray) = x
 aos_to_soa(x::TracedRArray) = x
 aos_to_soa(x::AnyTracedRArray) = x
 
+aos_to_soa(x::Array{Union{}}) = x
+
 function aos_to_soa(x::Array{<:TracedRNumber{T}}) where {T}
     isa_traced_soa(ancestor(x)) && return x
     for i in eachindex(x)
@@ -230,6 +232,9 @@ use_overlayed_version(::Number) = false
 use_overlayed_version(::MissingTracedValue) = true
 use_overlayed_version(rng::ReactantRNG) = use_overlayed_version(rng.seed)
 use_overlayed_version(::AbstractArray{<:TracedRNumber}) = true
+# `Union{}`-eltype (empty) arrays match the covariant traced-array patterns but
+# hold no traced values
+use_overlayed_version(::AbstractArray{Union{}}) = false
 use_overlayed_version(::TracedRArray) = true
 use_overlayed_version(::TracedRNumber) = true
 use_overlayed_version(::TracedStepRangeLen) = true
diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index 855116d0f6..ca7107ac28 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -312,13 +312,17 @@ for (jlop, hloop) in (
     (:(Base.:*), :multiply),
     (:(Base.:/), :divide),
     (:(Base.:^), :power),
-    (:(Base.rem), :remainder),
 )
     @eval function $(jlop)(lhs::TracedRNumber{T}, rhs::TracedRNumber{T}) where {T}
         return @opcall $(hloop)(lhs, rhs)
     end
 end
 
+# real-only, so that it dominates the mixed-argument methods below
+function Base.rem(lhs::TracedRReal{T}, rhs::TracedRReal{T}) where {T}
+    return @opcall remainder(lhs, rhs)
+end
+
 function Base.rem(x::TracedRReal, y::TracedRReal, ::typeof(Base.RoundFromZero))
     return ifelse(
         signbit(x) == signbit(y), rem(x, y, Base.RoundUp), rem(x, y, Base.RoundDown)
@@ -336,14 +340,14 @@ Base.:*(z::Complex{Bool}, x::TracedRReal) = x * z
 
 # Based on https://github.com/JuliaLang/julia/blob/39255d47db7657950ff1c82137ecec5a70bae622/base/float.jl#L608-L617
 function Base.mod(
-    @nospecialize(x::Reactant.TracedRNumber{T}), @nospecialize(y::Reactant.TracedRNumber{T})
+    @nospecialize(x::TracedRReal{T}), @nospecialize(y::TracedRReal{T})
 ) where {T}
     r = rem(x, y)
     return ifelse(r == 0, copysign(r, y), ifelse((r > 0) ⊻ (y > 0), r + y, r))
 end
 
 function Base.mod1(
-    @nospecialize(x::Reactant.TracedRNumber{T}), @nospecialize(y::Reactant.TracedRNumber{T})
+    @nospecialize(x::TracedRReal{T}), @nospecialize(y::TracedRReal{T})
 ) where {T}
     m = mod(x, y)
     return ifelse(m == 0, y, m)
@@ -982,6 +986,15 @@ function Base.cis(x::TracedRReal)
     return complex(c, s)
 end
 
+Base.sincos(x::TracedRFloat) = (sin(x), cos(x))
+Base.exp2(x::TracedRFloat{T}) where {T} = T(2)^x
+Base.exp10(x::TracedRFloat{T}) where {T} = T(10)^x
+Base.mod2pi(x::TracedRFloat{T}) where {T} = mod(x, T(2π))
+function Base.modf(x::TracedRFloat)
+    ipart = trunc(x)
+    return (x - ipart, ipart)
+end
+
 function Base.sinc(x::TracedRNumber)
     r = ifelse(iszero(x), one(x), sinpi(x) / (pi * x))
     return ifelse(isinf(x), zero(x), r)
@@ -1004,6 +1017,8 @@ for TracedT in (:TracedRReal, :TracedRComplex)
 end
 Base.Math._log(x::TracedRNumber, base, ::Symbol) = log(x) / log(Reactant._unwrap_val(base))
 
+Base.isreal(x::TracedRComplex) = iszero(imag(x))
+
 Base.isodd(x::TracedRComplex) = isodd(real(x))
 Base.isodd(x::TracedRInteger) = !iszero(rem(x, 2))
 function Base.isodd(x::TracedRFloat)
diff --git a/src/TracedRational.jl b/src/TracedRational.jl
index a540367a28..d1cfc59734 100644
--- a/src/TracedRational.jl
+++ b/src/TracedRational.jl
@@ -93,6 +93,16 @@ end
 function Base.promote_rule(::Type{BigFloat}, ::Type{Reactant.TracedRational{T}}) where {T}
     return BigFloat
 end
+function Base.promote_rule(
+    ::Type{<:Reactant.TracedRFloat{T2}}, ::Type{TracedRational{T}}
+) where {T2,T}
+    return Reactant.TracedRFloat{T2}
+end
+function Base.promote_rule(
+    ::Type{TracedRational{T}}, ::Type{<:Reactant.TracedRFloat{T2}}
+) where {T2,T}
+    return Reactant.TracedRFloat{T2}
+end
 
 # Operations
 Base.sign(x::TracedRational) = oftype(x, sign(x.num))
diff --git a/src/TracedUtils.jl b/src/TracedUtils.jl
index a912f57664..596af5f5a6 100644
--- a/src/TracedUtils.jl
+++ b/src/TracedUtils.jl
@@ -81,6 +81,10 @@ function ReactantCore.materialize_traced_array(
     return permutedims(materialize_traced_array(parent(x)), perm)
 end
 
+# Guards against `Union{}`-eltype (empty) arrays, which match the covariant
+# `AnyTracedRArray` alias but carry no traced values.
+ReactantCore.materialize_traced_array(x::AbstractArray{Union{}}) = x
+
 function ReactantCore.materialize_traced_array(
     x::AbstractArray{<:TracedRNumber{T}}
 ) where {T}
diff --git a/src/Tracing.jl b/src/Tracing.jl
index 7fc502826c..0a41dc7ff5 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -989,10 +989,18 @@ returns
 
 The first type parameter has been promoted to satisfy to be in agreement with the second parameter.
 """
+# `TracedRArray{T,N}` with the element-type parameter left free has no single
+# eltype and hence cannot instantiate an `AbstractArray{T}`-bounded type
+# parameter; substitute the concrete traced array type.
+concretize_traced_param(@nospecialize(T)) = T
+function concretize_traced_param(::Type{TracedRArray{T,N}}) where {T,N}
+    return TracedRArray{T,N,traced_number_type(T)}
+end
+
 function apply_type_with_promotion(wrapper, params, relevant_typevars=typevar_dict(wrapper))
     unwrapped = Base.unwrap_unionall(wrapper) # remove all the typevars
     original_params = copy(params)
-    params = [params...]
+    params = Any[concretize_traced_param(p) for p in params]
 
     changed = true
     iter = 0
diff --git a/src/Types.jl b/src/Types.jl
index 0d34dec85a..4ed61f2f19 100644
--- a/src/Types.jl
+++ b/src/Types.jl
@@ -270,6 +270,17 @@ function ConcretePJRTNumber(data::T; kwargs...) where {T<:Number}
     return ConcretePJRTNumber{T}(data; kwargs...)
 end
 
+# Value constructors for the kind-specific types, e.g. via
+# `convert(ConcretePJRTFloat{Float64,1}, x)`. The `Rational`/`Complex` methods
+# disambiguate against Base's constructors on `Integer`/`AbstractFloat`/`Real`.
+for CT in (:ConcretePJRTInteger, :ConcretePJRTFloat, :ConcretePJRTComplex),
+    XT in (:Number, :Rational, :Complex)
+
+    @eval function (::Type{CT})(x::$XT; kwargs...) where {CT<:$CT}
+        return ConcretePJRTNumber{unwrapped_eltype(CT)}(x; kwargs...)
+    end
+end
+
 function ConcretePJRTNumber(data::ConcretePJRTNumber; kwargs...)
     return ConcretePJRTNumber(
         to_number(data);
@@ -431,6 +442,14 @@ function ConcreteIFRTNumber(data::T; kwargs...) where {T<:Number}
     return ConcreteIFRTNumber{T}(data; kwargs...)
 end
 
+for CT in (:ConcreteIFRTInteger, :ConcreteIFRTFloat, :ConcreteIFRTComplex),
+    XT in (:Number, :Rational, :Complex)
+
+    @eval function (::Type{CT})(x::$XT; kwargs...) where {CT<:$CT}
+        return ConcreteIFRTNumber{unwrapped_eltype(CT)}(x; kwargs...)
+    end
+end
+
 function ConcreteIFRTNumber(data::ConcreteIFRTNumber; kwargs...)
     return ConcreteIFRTNumber(
         to_number(data);
diff --git a/test/core/tracing.jl b/test/core/tracing.jl
index 4a6c1782d4..32a4b63ec7 100644
--- a/test/core/tracing.jl
+++ b/test/core/tracing.jl
@@ -26,7 +26,11 @@ struct RMSProp{Teta,Trho,Teps,C<:Bool}
 end
 
 @testset "Traced Type" begin
-    @test !(Vector{Union{}} <: Reactant.AnyTracedRArray)
+    # `Vector{Union{}}` (e.g. an empty array literal) matches the covariant
+    # `AnyTracedRArray` alias, but must not be routed into the traced-array
+    # machinery
+    @test ReactantCore.materialize_traced_array(Union{}[]) isa Vector{Union{}}
+    @test Reactant.aos_to_soa(Union{}[]) isa Vector{Union{}}
 end
 
 mul(a, b) = a .* b
@@ -338,11 +342,11 @@ end
         Foo, [Float64, Bar{Float64}, Reactant.TracedRArray{Float64,1}]
     ) == (
         Foo{
-            TracedRNumber{Float64},
-            Bar{TracedRNumber{Float64}},
-            Reactant.TracedRArray{Float64,1},
+            Reactant.TracedRFloat{Float64},
+            Bar{Reactant.TracedRFloat{Float64}},
+            Reactant.TracedRArray{Float64,1,Reactant.TracedRFloat{Float64}},
         },
-        [true, true, false],
+        [true, true, true],
     )
 
     @test Reactant.apply_type_with_promotion(

From f31642af2e3153885e0cce8d68d66d79962f07f2 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Tue, 7 Jul 2026 23:42:56 +0200
Subject: [PATCH 11/30] Fix traced-integer exponents, Active scalar autodiff,
 and type concretization

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.
---
 src/Enzyme.jl        | 38 ++++++++++++++++++++++++++++----------
 src/Indexing.jl      | 11 +++++++++++
 src/TracedRArray.jl  |  9 +++++++--
 src/TracedRNumber.jl | 16 ++++++++++++++++
 src/Tracing.jl       |  4 ++--
 test/core/compile.jl |  4 +++-
 6 files changed, 67 insertions(+), 15 deletions(-)

diff --git a/src/Enzyme.jl b/src/Enzyme.jl
index 987f68dc02..e10b694ebb 100644
--- a/src/Enzyme.jl
+++ b/src/Enzyme.jl
@@ -67,6 +67,12 @@ for RT in (:RInteger, :RFloat, :RComplex)
     @eval @inline Enzyme.make_zero(x::$RT) = zero(Core.Typeof(x))
 end
 
+# Reactant numbers are mutable wrappers around device/MLIR values and must not
+# be classified as active scalars, which Enzyme does for any `AbstractFloat`.
+function Enzyme.Compiler.active_reg_nothrow(::Type{T}) where {T<:RNumber}
+    return Enzyme.Compiler.DupState
+end
+
 # Reactant numbers are mutable wrappers around device/MLIR values and must
 # stay `Duplicated` under autodiff even though they subtype `AbstractFloat`.
 for RT in (:RArray, :RNumber, :RInteger, :RFloat, :RComplex)
@@ -468,7 +474,11 @@ function overload_autodiff(
                 push!(ret_activity, act)
 
                 if act == EnzymeActivity.OUT || act == EnzymeActivity.OUT_NO_NEED
-                    seed = reverse_seeds[path]
+                    seed = get(reverse_seeds, path) do
+                        # `Active` arguments have no shadow; the re-emitted
+                        # argument result receives no incoming cotangent.
+                        TracedUtils.get_mlir_data(zero(a))
+                    end
                     push!(ad_inputs, seed)
                 end
             else
@@ -579,7 +589,7 @@ function overload_autodiff(
         end
     end
 
-    restup = Any[(a isa Active) ? copy(a) : nothing for a in args]
+    restup = Any[nothing for _ in args]
     for a in linear_args
         idx, path = TracedUtils.get_argidx(a, argprefix)
 
@@ -590,14 +600,22 @@ function overload_autodiff(
             @assert false
         end
 
-        set_act!(
-            arg,
-            path[3:end],
-            reverse,
-            TracedUtils.transpose_val(MLIR.IR.result(res, residx));
-            width,
-            emptypath=arg isa Active,
-        )
+        if arg isa Active && width == 1 && isempty(path[3:end])
+            # `Active` arguments return their gradient in the result tuple
+            # instead of accumulating into a shadow
+            restup[idx - fnwrap] = Reactant.TracedRNumber{unwrapped_eltype(arg.val)}(
+                (), TracedUtils.transpose_val(MLIR.IR.result(res, residx))
+            )
+        else
+            set_act!(
+                arg,
+                path[3:end],
+                reverse,
+                TracedUtils.transpose_val(MLIR.IR.result(res, residx));
+                width,
+                emptypath=arg isa Active,
+            )
+        end
         residx += 1
     end
 
diff --git a/src/Indexing.jl b/src/Indexing.jl
index e2b6a7a4d7..0461f8814f 100644
--- a/src/Indexing.jl
+++ b/src/Indexing.jl
@@ -214,6 +214,17 @@ function Base.getindex(
     return out
 end
 
+# Disambiguate the TracedRange getindex methods against the range loops above
+function Base.getindex(v::TracedUnitRange{<:TracedRNumber}, i::TracedRInteger)
+    return convert(eltype(v), v.start + (i - oneunit(i)))
+end
+function Base.getindex(v::TracedUnitRange{<:TracedRNumber}, i::TracedRInteger{Int})
+    return convert(eltype(v), v.start + (i - oneunit(i)))
+end
+function Base.getindex(r::TracedStepRangeLen{<:TracedRNumber}, i::TracedRInteger{Int})
+    return Base.unsafe_getindex(r, i)
+end
+
 ## StepRangeLen Indexing
 function Base.getindex(r::TracedStepRangeLen{T}, s::OrdinalRange{S}) where {T,S}
     @inline
diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
index ad12a279e0..fe6b6e1ff5 100644
--- a/src/TracedRArray.jl
+++ b/src/TracedRArray.jl
@@ -304,6 +304,10 @@ function Base.similar(
     return (@opcall fill(zero(T), dims))::TracedRArray{T,N}
 end
 
+broadcast_eltype_param(::Type{<:TracedRNumber{T}}) where {T} = T
+broadcast_eltype_param(::Type{T}) where {T<:Reactant.ReactantPrimitive} = T
+broadcast_eltype_param(@nospecialize(_)) = nothing
+
 function Base.copy(bc::Broadcasted{<:AbstractReactantArrayStyle{0}})
     ElType = Broadcast.combine_eltypes(bc.f, bc.args)
     dest = copyto!(similar(bc, ElType), bc)
@@ -324,8 +328,9 @@ function _copy(bc)
         bc.f
     end
     ElType = Broadcast.combine_eltypes(fn, bc.args)
-    # Special case a union{} return so we can see the better error message
-    if ElType === Union{} || ElType == Any || ElType == TracedRNumber
+    # Inference may return an abstract eltype (e.g. `Real` or `Union{}`); get a
+    # usable one from a sample value in that case.
+    if broadcast_eltype_param(ElType) === nothing
         ElType = Core.Typeof(fn(map(first_scalar, bc.args)...))
     end
     if ElType == Any || ElType == Union{}
diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index ca7107ac28..b62d14138e 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -13,6 +13,7 @@ using ..Reactant:
     traced_number_type
 using ..Ops: @opcall
 using ReactantCore: ReactantCore, @trace
+using BFloat16s: BFloat16
 using Adapt: Adapt
 
 # This isn't technically necessary in this module, but this type used to be
@@ -858,6 +859,21 @@ for (jlop, hloop) in ((:(Base.:&), :and), (:(Base.:|), :or), (:(Base.xor), :xor)
 end
 Base.:!(x::TracedRInteger) = @opcall not(x)
 
+# With a traced integer exponent, Base's `^(::Number, ::Integer)` would call
+# `power_by_squaring`, which branches on traced booleans. The extra methods
+# disambiguate against Base's specialized `^` methods.
+for B in (:TracedRInteger, :TracedRFloat, :TracedRComplex, :Real, :Complex)
+    @eval Base.:^(x::$B, p::TracedRInteger) = ^(promote(x, p)...)
+end
+for B in (
+    Float16, Float32, Union{Float16,Float32}, Float64, BFloat16,
+    Complex{<:AbstractFloat}, Complex{<:Integer}, Complex{<:Rational},
+)
+    @eval Base.:^(x::$B, p::TracedRInteger) = ^(promote(x, p)...)
+end
+
+Base.fma(x::TracedRFloat{T}, y::TracedRFloat{T}, z::TracedRFloat{T}) where {T} = x * y + z
+
 function Base.literal_pow(
     ::Base.RefValue{typeof(^)}, x::TracedRNumber{T}, ::Base.RefValue{Val{P}}
 ) where {T,P}
diff --git a/src/Tracing.jl b/src/Tracing.jl
index 0a41dc7ff5..8983651318 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -302,7 +302,7 @@ Base.@nospecializeinfer function traced_type_inner(
     if mode == ConcreteToTraced
         A´ = Base.unwrap_unionall(A)
         T, N, _ = A´.parameters
-        A´´ = TracedRArray{T,N}
+        A´´ = T isa Core.TypeVar ? TracedRArray{T,N} : TracedRArray{T,N,traced_number_type(T)}
         A_ret = N isa Core.TypeVar ? UnionAll(N, A´´) : A´´
         A_ret2 = T isa Core.TypeVar ? UnionAll(T, A_ret) : A_ret
         return A_ret2
@@ -332,7 +332,7 @@ Base.@nospecializeinfer function traced_type_inner(
     if mode == ConcreteToTraced
         A´ = Base.unwrap_unionall(A)
         T, N = A´.parameters
-        A´´ = TracedRArray{T,N}
+        A´´ = T isa Core.TypeVar ? TracedRArray{T,N} : TracedRArray{T,N,traced_number_type(T)}
         A_ret = N isa Core.TypeVar ? UnionAll(N, A´´) : A´´
         A_ret2 = T isa Core.TypeVar ? UnionAll(T, A_ret) : A_ret
         return A_ret2
diff --git a/test/core/compile.jl b/test/core/compile.jl
index 9afd7435a3..5eddd4eeab 100644
--- a/test/core/compile.jl
+++ b/test/core/compile.jl
@@ -10,7 +10,9 @@ intout_caller(vis) = @noinline intout(vis)
 @testset "compile" begin
     vis = rand(Float64, 64)
     visr = Reactant.to_rarray(vis)
-    @test_throws MethodError @compile intout_caller(visr)
+    # Traced arrays dispatch as `AbstractArray{<:Real}` (#1570)
+    fn = @compile intout_caller(visr)
+    @test size(fn(visr)) == size(vis)
 end
 
 @testset "create_result" begin

From bd1ef8e1fdb6b87fb8988766f92a3e576b0f09d4 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Wed, 8 Jul 2026 00:28:45 +0200
Subject: [PATCH 12/30] Add number-kind hierarchy tests and format

Created by generative AI.
---
 src/ConcreteRArray.jl     | 12 ++----
 src/Enzyme.jl             |  2 +-
 src/TracedRArray.jl       |  7 +---
 src/TracedRNumber.jl      | 18 +++++----
 src/Tracing.jl            |  6 ++-
 test/core/compile.jl      |  9 ++---
 test/core/type_numbers.jl | 78 +++++++++++++++++++++++++++++++++++++++
 7 files changed, 101 insertions(+), 31 deletions(-)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 97483e4c11..4b444abc7a 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -300,10 +300,8 @@ const BASE_SPECIFIC_ROUNDING_MODES = (
 
 for T1 in CONCRETE_NUMBER_KINDS
     @eval begin
-        Base.div(x::$(T1), y::Number, r::RoundingMode=RoundToZero) =
-            div(to_number(x), y, r)
-        Base.div(x::Number, y::$(T1), r::RoundingMode=RoundToZero) =
-            div(x, to_number(y), r)
+        Base.div(x::$(T1), y::Number, r::RoundingMode=RoundToZero) = div(to_number(x), y, r)
+        Base.div(x::Number, y::$(T1), r::RoundingMode=RoundToZero) = div(x, to_number(y), r)
         Base.div(x::$(T1), y::Real, r::RoundingMode=RoundToZero) = div(to_number(x), y, r)
         Base.div(x::Real, y::$(T1), r::RoundingMode=RoundToZero) = div(x, to_number(y), r)
         Base.div(x::$(T1), y::Rational, r::RoundingMode=RoundToZero) =
@@ -327,8 +325,7 @@ for T1 in CONCRETE_NUMBER_KINDS
         @eval Base.div(x::$(T1), y::$(T2), r::RoundingMode=RoundToZero) =
             div(to_number(x), to_number(y), r)
         for RM in BASE_SPECIFIC_ROUNDING_MODES
-            @eval Base.div(x::$(T1), y::$(T2), r::$RM) =
-                div(to_number(x), to_number(y), r)
+            @eval Base.div(x::$(T1), y::$(T2), r::$RM) = div(to_number(x), to_number(y), r)
         end
     end
     for T2 in TRACED_NUMBER_KINDS
@@ -380,8 +377,7 @@ for jlop in (:(Base.isnan), :(Base.isfinite)),
 end
 
 isapprox_pairs = Any[
-    (AbstractConcreteArray{<:Any,0}, Number),
-    (Number, AbstractConcreteArray{<:Any,0}),
+    (AbstractConcreteArray{<:Any,0}, Number), (Number, AbstractConcreteArray{<:Any,0})
 ]
 for T1 in CONCRETE_NUMBER_KINDS
     push!(isapprox_pairs, (T1, Number), (Number, T1), (T1, Integer), (Integer, T1))
diff --git a/src/Enzyme.jl b/src/Enzyme.jl
index e10b694ebb..60162295f5 100644
--- a/src/Enzyme.jl
+++ b/src/Enzyme.jl
@@ -477,7 +477,7 @@ function overload_autodiff(
                     seed = get(reverse_seeds, path) do
                         # `Active` arguments have no shadow; the re-emitted
                         # argument result receives no incoming cotangent.
-                        TracedUtils.get_mlir_data(zero(a))
+                        return TracedUtils.get_mlir_data(zero(a))
                     end
                     push!(ad_inputs, seed)
                 end
diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
index fe6b6e1ff5..29d2379f12 100644
--- a/src/TracedRArray.jl
+++ b/src/TracedRArray.jl
@@ -4,12 +4,7 @@ using Base: Broadcast
 using Base.Broadcast: Broadcasted, AbstractArrayStyle, instantiate
 
 using ..Reactant:
-    Reactant,
-    TracedRArray,
-    TracedRNumber,
-    TracedRInteger,
-    AnyTracedRArray,
-    AnyTracedRVector
+    Reactant, TracedRArray, TracedRNumber, TracedRInteger, AnyTracedRArray, AnyTracedRVector
 using ..Reactant: MLIR, unwrapped_eltype
 using ..Ops: @opcall
 using ..TracedUtils: TracedUtils, get_mlir_data, set_mlir_data!, materialize_traced_array
diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index b62d14138e..bbf0b8df21 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -161,9 +161,7 @@ function Base.promote_rule(
 ) where {T<:Integer,S}
     return traced_number_type(Base.promote_type(Rational{T}, S))
 end
-function Base.promote_rule(
-    ::Type{Complex{T}}, ::Type{<:TracedRReal{S}}
-) where {T<:Real,S}
+function Base.promote_rule(::Type{Complex{T}}, ::Type{<:TracedRReal{S}}) where {T<:Real,S}
     return traced_number_type(Base.promote_type(Complex{T}, S))
 end
 
@@ -817,9 +815,7 @@ function Base.:+(
     return Base.TwicePrecision(Base.canonicalize2(r, s)...)
 end
 
-function Base.:*(
-    x::TwicePrecision{<:Union{Float16,Float32,Float64}}, v::TracedRInteger
-)
+function Base.:*(x::TwicePrecision{<:Union{Float16,Float32,Float64}}, v::TracedRInteger)
     return invoke(*, Tuple{TwicePrecision,TracedRNumber}, x, v)
 end
 
@@ -866,8 +862,14 @@ for B in (:TracedRInteger, :TracedRFloat, :TracedRComplex, :Real, :Complex)
     @eval Base.:^(x::$B, p::TracedRInteger) = ^(promote(x, p)...)
 end
 for B in (
-    Float16, Float32, Union{Float16,Float32}, Float64, BFloat16,
-    Complex{<:AbstractFloat}, Complex{<:Integer}, Complex{<:Rational},
+    Float16,
+    Float32,
+    Union{Float16,Float32},
+    Float64,
+    BFloat16,
+    Complex{<:AbstractFloat},
+    Complex{<:Integer},
+    Complex{<:Rational},
 )
     @eval Base.:^(x::$B, p::TracedRInteger) = ^(promote(x, p)...)
 end
diff --git a/src/Tracing.jl b/src/Tracing.jl
index 8983651318..f9da34f986 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -302,7 +302,8 @@ Base.@nospecializeinfer function traced_type_inner(
     if mode == ConcreteToTraced
         A´ = Base.unwrap_unionall(A)
         T, N, _ = A´.parameters
-        A´´ = T isa Core.TypeVar ? TracedRArray{T,N} : TracedRArray{T,N,traced_number_type(T)}
+        A´´ =
+            T isa Core.TypeVar ? TracedRArray{T,N} : TracedRArray{T,N,traced_number_type(T)}
         A_ret = N isa Core.TypeVar ? UnionAll(N, A´´) : A´´
         A_ret2 = T isa Core.TypeVar ? UnionAll(T, A_ret) : A_ret
         return A_ret2
@@ -332,7 +333,8 @@ Base.@nospecializeinfer function traced_type_inner(
     if mode == ConcreteToTraced
         A´ = Base.unwrap_unionall(A)
         T, N = A´.parameters
-        A´´ = T isa Core.TypeVar ? TracedRArray{T,N} : TracedRArray{T,N,traced_number_type(T)}
+        A´´ =
+            T isa Core.TypeVar ? TracedRArray{T,N} : TracedRArray{T,N,traced_number_type(T)}
         A_ret = N isa Core.TypeVar ? UnionAll(N, A´´) : A´´
         A_ret2 = T isa Core.TypeVar ? UnionAll(T, A_ret) : A_ret
         return A_ret2
diff --git a/test/core/compile.jl b/test/core/compile.jl
index 5eddd4eeab..da6fd20cf2 100644
--- a/test/core/compile.jl
+++ b/test/core/compile.jl
@@ -236,12 +236,9 @@ end
             sym = Symbol("result", var_idx[])
             var_idx[] += 1
 
-            push!(
-                resultgen_code,
-                quote
-                    $sym = ($MockTestCustomPath)($ar)
-                end,
-            )
+            push!(resultgen_code, quote
+                $sym = ($MockTestCustomPath)($ar)
+            end)
             result_cache[tocopy] = sym
         end
 
diff --git a/test/core/type_numbers.jl b/test/core/type_numbers.jl
index 7ca2673e44..858766bb76 100644
--- a/test/core/type_numbers.jl
+++ b/test/core/type_numbers.jl
@@ -276,3 +276,81 @@ end
 
     @test res ≈ tobc((10.0, 10.0), x)
 end
+
+@testset "Number type hierarchy (#1570)" begin
+    using Reactant:
+        RInteger,
+        RFloat,
+        RComplex,
+        RReal,
+        RNumber,
+        TracedRNumber,
+        TracedRReal,
+        TracedRInteger,
+        TracedRFloat,
+        TracedRComplex,
+        traced_number_type
+
+    @testset "abstract kinds" begin
+        @test RInteger <: Integer
+        @test RFloat <: AbstractFloat
+        @test RComplex <: Number
+        @test RReal <: Real
+    end
+
+    @testset "traced types" begin
+        @test TracedRInteger{Int64} <: RInteger{Int64} <: Integer
+        @test TracedRInteger{Bool} <: Integer
+        @test TracedRFloat{Float64} <: RFloat{Float64} <: AbstractFloat
+        @test TracedRComplex{ComplexF64} <: RComplex{ComplexF64} <: Number
+        @test TracedRReal <: Real
+        @test TracedRFloat{Float32} <: TracedRNumber{Float32}
+        @test traced_number_type(Int32) === TracedRInteger{Int32}
+        @test traced_number_type(Bool) === TracedRInteger{Bool}
+        @test traced_number_type(Float16) === TracedRFloat{Float16}
+        @test traced_number_type(ComplexF32) === TracedRComplex{ComplexF32}
+    end
+
+    @testset "concrete types" begin
+        @test ConcretePJRTInteger{Int64,1} <:
+            Reactant.AbstractConcreteInteger{Int64} <:
+            Integer
+        @test ConcretePJRTFloat{Float64,1} <:
+            Reactant.AbstractConcreteFloat{Float64} <:
+            AbstractFloat
+        @test ConcretePJRTComplex{ComplexF64,1} <: Number
+        @test ConcreteIFRTInteger{Int64} <: Integer
+        @test ConcreteIFRTFloat{Float64} <: AbstractFloat
+        @test ConcreteIFRTNumber{Float64} isa Type
+    end
+
+    @testset "traced arrays expose the numeric kind" begin
+        RT = traced_number_type(Float64)
+        @test Reactant.TracedRArray{Float64,2,RT} <: AbstractArray{<:AbstractFloat,2}
+        @test Reactant.TracedRArray{Float64,2,RT} <: AbstractArray{<:Real,2}
+        @test eltype(Reactant.TracedRArray{Int32,1,traced_number_type(Int32)}) <: Integer
+    end
+
+    @testset "promotion" begin
+        @test promote_type(TracedRFloat{Float64}, TracedRInteger{Int64}) ===
+            TracedRFloat{Float64}
+        @test promote_type(TracedRInteger{Int32}, Float64) === TracedRFloat{Float64}
+        @test promote_type(Bool, TracedRFloat{Float32}) === TracedRFloat{Float32}
+        @test promote_type(typeof(pi), TracedRFloat{Float32}) === TracedRFloat{Float32}
+        @test promote_type(TracedRReal{Float64}, TracedRComplex{ComplexF32}) ===
+            TracedRComplex{ComplexF64}
+    end
+
+    @testset "dispatch on Real/Integer while tracing" begin
+        real_only(x::Real) = 2x + 1
+        int_only(n::Integer) = n * n
+        x = Reactant.to_rarray([1.0, 2.0, 3.0])
+        res = @jit (v -> @allowscalar(real_only(v[1]) + real_only(v[2])))(x)
+        @test res isa AbstractFloat
+        @test res ≈ 8.0
+        xi = Reactant.to_rarray([2, 3, 4])
+        resi = @jit (v -> @allowscalar(int_only(v[1])))(xi)
+        @test resi isa Integer
+        @test resi == 4
+    end
+end

From 1c65ac4bf7523a93c4b07e16b0656e1129e47c2a Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Wed, 8 Jul 2026 07:50:16 +0200
Subject: [PATCH 13/30] Split CuTracedRNumber by numeric kind in the CUDA
 extension

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.
---
 ext/ReactantCUDAExt.jl   | 126 ++++++++++++++++++++++++++++-----------
 test/integration/cuda.jl |   2 +-
 2 files changed, 93 insertions(+), 35 deletions(-)

diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl
index dabadd2bfb..603860732f 100644
--- a/ext/ReactantCUDAExt.jl
+++ b/ext/ReactantCUDAExt.jl
@@ -63,27 +63,53 @@ end
 
 Reactant.use_overlayed_version(::CuTracedArray) = true
 
-struct CuTracedRNumber{T,A} <: Number
-    ptr::Core.LLVMPtr{T,A}
+# Split by numeric kind, mirroring the traced host types, so that kernel-side
+# numbers subtype Integer and AbstractFloat where possible.
+for (CuT, AbstractT) in (
+    (:CuTracedRInteger, :(Reactant.RInteger)),
+    (:CuTracedRFloat, :(Reactant.RFloat)),
+    (:CuTracedRComplex, :(Reactant.RComplex)),
+)
+    @eval begin
+        struct $CuT{T,A} <: $AbstractT{T}
+            ptr::Core.LLVMPtr{T,A}
 
-    function CuTracedRNumber{T,A}(xs::TracedRNumber) where {T,A}
-        Reactant.Compiler.guard_from_gc_for_module(MLIR.IR.current_module(), xs)
-        ptr = Base.reinterpret(Core.LLVMPtr{T,CUDA.AS.Global}, Base.pointer_from_objref(xs))
-        return new(ptr)
-    end
-    function CuTracedRNumber{T,A}(ptr::Core.LLVMPtr{T,A}) where {T,A}
-        return new(ptr)
+            function $CuT{T,A}(xs::TracedRNumber) where {T,A}
+                Reactant.Compiler.guard_from_gc_for_module(MLIR.IR.current_module(), xs)
+                ptr = Base.reinterpret(
+                    Core.LLVMPtr{T,CUDA.AS.Global}, Base.pointer_from_objref(xs)
+                )
+                return new(ptr)
+            end
+            function $CuT{T,A}(ptr::Core.LLVMPtr{T,A}) where {T,A}
+                return new(ptr)
+            end
+        end
+
+        Reactant.use_overlayed_version(::$CuT) = true
     end
 end
 
-Reactant.use_overlayed_version(::CuTracedRNumber) = true
+const CuTracedRReal{T,A} = Union{CuTracedRInteger{T,A},CuTracedRFloat{T,A}}
+const CuTracedRNumber{T,A} = Union{CuTracedRReal{T,A},CuTracedRComplex{T,A}}
+
+@inline cu_traced_number_type(::Type{T}) where {T<:Integer} = CuTracedRInteger{T}
+@inline cu_traced_number_type(::Type{T}) where {T<:AbstractFloat} = CuTracedRFloat{T}
+@inline cu_traced_number_type(::Type{T}) where {T<:Complex} = CuTracedRComplex{T}
+
+@inline function CuTracedRNumber{T,A}(xs::TracedRNumber) where {T,A}
+    return cu_traced_number_type(T){A}(xs)
+end
+@inline function CuTracedRNumber{T,A}(ptr::Core.LLVMPtr{T,A}) where {T,A}
+    return cu_traced_number_type(T){A}(ptr)
+end
 
 Base.@nospecializeinfer Reactant.is_traced_number(
     @nospecialize(T::Type{<:CuTracedRNumber})
 ) = true
 Reactant.unwrapped_eltype(::Type{<:CuTracedRNumber{T}}) where {T} = T
 
-@inline CuTracedRNumber{T,A}(val::Number) where {T,A} = convert(CuTracedRNumber{T,A}, val)
+@inline (::Type{CuT})(val::Number) where {CuT<:CuTracedRNumber} = convert(CuT, val)
 
 function Base.getindex(RN::CuTracedRNumber{T,A}) where {T,A}
     align = alignment(RN)
@@ -92,6 +118,10 @@ end
 
 Base.convert(::Type{T}, RN::CuTracedRNumber) where {T<:Number} = convert(T, getindex(RN))
 
+# Defined per numeric kind: a single method on the `CuTracedRNumber` union
+# would be ambiguous with Base's methods on `Real`/`Integer`/`AbstractFloat`.
+const CU_TRACED_NUMBER_KINDS = (CuTracedRInteger, CuTracedRFloat, CuTracedRComplex)
+
 for jlop in (
     :(Base.min),
     :(Base.mod),
@@ -106,10 +136,32 @@ for jlop in (
     :(Base.:(==)),
     :(Base.:(!=)),
 )
-    @eval begin
-        @inline $jlop(a::CuTracedRNumber, b::CuTracedRNumber) = $jlop(a[], b[])
-        @inline $jlop(a::CuTracedRNumber{T,A}, b::Number) where {T,A} = $jlop(a[], b)
-        @inline $jlop(a::Number, b::CuTracedRNumber{T,A}) where {T,A} = $jlop(a, b[])
+    for K1 in CU_TRACED_NUMBER_KINDS
+        @eval begin
+            @inline $jlop(a::$K1, b::Number) = $jlop(a[], b)
+            @inline $jlop(a::Number, b::$K1) = $jlop(a, b[])
+        end
+        for K2 in CU_TRACED_NUMBER_KINDS
+            @eval @inline $jlop(a::$K1, b::$K2) = $jlop(a[], b[])
+        end
+        for X in (
+            :Real,
+            :Integer,
+            :Bool,
+            :AbstractFloat,
+            :Complex,
+            :(Complex{Bool}),
+            :Rational,
+            :BigInt,
+            :BigFloat,
+            :AbstractIrrational,
+        )
+            jlop == :(Base.:^) && X === :Integer && continue
+            @eval begin
+                @inline $jlop(a::$X, b::$K1) = $jlop(a, b[])
+                @inline $jlop(a::$K1, b::$X) = $jlop(a[], b)
+            end
+        end
     end
 end
 
@@ -123,20 +175,20 @@ end
 @inline Base.ifelse(cond::CuTracedRNumber, a::CuTracedRNumber, b::CuTracedRNumber) =
     ifelse(cond[], a[], b[])
 
-Base.@constprop :aggressive @inline Base.:^(
-    a::CuTracedRNumber{T,A}, b::Integer
-) where {T,A} = ^(a[], b)
+for K in CU_TRACED_NUMBER_KINDS
+    @eval Base.@constprop :aggressive @inline Base.:^(a::$K, b::Integer) = ^(a[], b)
+end
 
 @inline Base.unsafe_trunc(::Type{T}, a::CuTracedRNumber) where {T} =
     Base.unsafe_trunc(T, a[])
 
-for jlop in (:(Base.:+), :(Base.:-), :(Base.isnan), :(Base.isfinite), :(Base.isinf))
-    @eval begin
-        @inline $jlop(a::CuTracedRNumber) = $jlop(a[])
-    end
+for jlop in (:(Base.:+), :(Base.:-), :(Base.isnan), :(Base.isfinite), :(Base.isinf)),
+    K in CU_TRACED_NUMBER_KINDS
+
+    @eval @inline $jlop(a::$K) = $jlop(a[])
 end
 
-Base.OneTo(x::CuTracedRNumber{<:Integer}) = Base.OneTo(x[])
+Base.OneTo(x::CuTracedRInteger) = Base.OneTo(x[])
 
 @static if isdefined(Base, :unchecked_oneto)
     function Base.unchecked_oneto(x::CuTracedRNumber{<:Integer})
@@ -273,8 +325,12 @@ function Base.show(io::IO, a::AT) where {AT<:CuTracedArray}
     Printf.@printf(io, "%s cu traced array at %p", join(size(a), '×'), Int(pointer(a)))
 end
 
-function Base.show(io::IO, a::AT) where {AT<:CuTracedRNumber}
-    Printf.@printf(io, "%s cu traced rnumber at %p", join(size(a), '×'), Int(pointer(a)))
+for K in CU_TRACED_NUMBER_KINDS
+    @eval function Base.show(io::IO, a::AT) where {AT<:$K}
+        Printf.@printf(
+            io, "%s cu traced rnumber at %p", join(size(a), '×'), Int(pointer(a))
+        )
+    end
 end
 
 ## array interface
@@ -1597,15 +1653,17 @@ Base.@nospecializeinfer function Reactant.traced_type_inner(
     return A
 end
 
-Base.@nospecializeinfer function Reactant.traced_type_inner(
-    @nospecialize(A::Type{<:CuTracedRNumber}),
-    seen,
-    @nospecialize(mode::Reactant.TraceMode),
-    @nospecialize(track_numbers::Type),
-    @nospecialize(ndevices),
-    @nospecialize(runtime)
-)
-    return A
+for CuT in (:CuTracedRNumber, :CuTracedRInteger, :CuTracedRFloat, :CuTracedRComplex)
+    @eval Base.@nospecializeinfer function Reactant.traced_type_inner(
+        @nospecialize(A::Type{<:$CuT}),
+        seen,
+        @nospecialize(mode::Reactant.TraceMode),
+        @nospecialize(track_numbers::Type),
+        @nospecialize(ndevices),
+        @nospecialize(runtime)
+    )
+        return A
+    end
 end
 
 Base.@nospecializeinfer function Reactant.traced_type_inner(
diff --git a/test/integration/cuda.jl b/test/integration/cuda.jl
index f142c8babf..a6559fe932 100644
--- a/test/integration/cuda.jl
+++ b/test/integration/cuda.jl
@@ -273,7 +273,7 @@ end
     clock = MyClock(ConcreteRNumber(0.0), 3)
     params = RawParams(ConcreteRNumber(0.1), 7)
 
-    @test_throws "GPU kernel argument of type @NamedTuple{time::Reactant.TracedRNumber{Float64}, iteration::Int64} contains an unadapted traced value at field: time" Reactant.@compile raise =
+    @test_throws "GPU kernel argument of type @NamedTuple{time::Reactant.TracedRFloat{Float64}, iteration::Int64} contains an unadapted traced value at field: time" Reactant.@compile raise =
         true raise_first = true sync = true run!(arr, clock, params)
 
     clock2 = MyClock2(ConcreteRNumber(0.0), 3)

From 139bb91fd85f014a9de2a8dba7a32e0809979839 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Wed, 8 Jul 2026 07:50:16 +0200
Subject: [PATCH 14/30] Fix integration fallout and enforce the TracedRArray
 element invariant

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.
---
 ext/ReactantAbstractFFTsExt.jl     | 14 ++++++++++++++
 ext/ReactantTensorOperationsExt.jl |  3 ++-
 src/TracedRNumber.jl               |  2 +-
 src/Types.jl                       | 19 +++++++++++++++++--
 test/integration/structarrays.jl   | 10 ++++------
 5 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/ext/ReactantAbstractFFTsExt.jl b/ext/ReactantAbstractFFTsExt.jl
index 3cde0dd65a..417ee694cc 100644
--- a/ext/ReactantAbstractFFTsExt.jl
+++ b/ext/ReactantAbstractFFTsExt.jl
@@ -9,6 +9,20 @@ using Reactant.Ops: @opcall
 # catch its methods.
 const AnyStridedTracedRArray{T,N} = StridedArray{<:TracedRNumber{T},N}
 
+# Traced reals hit the `fftfloat`-based conversion path in AbstractFFTs, whose
+# fallback only supports the BLAS float types and copies through a host array.
+function AbstractFFTs._fftfloat(::Type{T}) where {T<:TracedRNumber}
+    return Reactant.traced_number_type(AbstractFFTs._fftfloat(Reactant.unwrapped_eltype(T)))
+end
+
+function AbstractFFTs.complexfloat(x::AbstractArray{<:Reactant.TracedRReal{T}}) where {T}
+    return Reactant.promote_to(Reactant.TracedRArray{complex(float(T)),ndims(x)}, x)
+end
+
+function AbstractFFTs.realfloat(x::AbstractArray{<:Reactant.TracedRReal{T}}) where {T}
+    return Reactant.promote_to(Reactant.TracedRArray{float(T),ndims(x)}, x)
+end
+
 # To automatically convert FFT plans to traced versions
 # To extend a user needs to extend Reactant.reactant_fftplan for their plan type
 # see ReactantFFTWExt.jl for an example implementation
diff --git a/ext/ReactantTensorOperationsExt.jl b/ext/ReactantTensorOperationsExt.jl
index 5e8de9c294..b37cd366b0 100644
--- a/ext/ReactantTensorOperationsExt.jl
+++ b/ext/ReactantTensorOperationsExt.jl
@@ -17,7 +17,8 @@ function TO.tensoradd_type(TC, A::ConcreteRArray, pA::Index2Tuple, conjA::Bool)
 end
 
 function TO.tensoradd_type(TC, A::TracedRArray, pA::Index2Tuple, conjA::Bool)
-    return TracedRArray{unwrapped_eltype(TC),TO.numind(pA)}
+    T = unwrapped_eltype(TC)
+    return TracedRArray{T,TO.numind(pA),Reactant.traced_number_type(T)}
 end
 
 # backend selection
diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index bbf0b8df21..7b0677ea76 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -858,7 +858,7 @@ Base.:!(x::TracedRInteger) = @opcall not(x)
 # With a traced integer exponent, Base's `^(::Number, ::Integer)` would call
 # `power_by_squaring`, which branches on traced booleans. The extra methods
 # disambiguate against Base's specialized `^` methods.
-for B in (:TracedRInteger, :TracedRFloat, :TracedRComplex, :Real, :Complex)
+for B in (:TracedRInteger, :TracedRFloat, :TracedRComplex, :Real, :Complex, :Rational)
     @eval Base.:^(x::$B, p::TracedRInteger) = ^(promote(x, p)...)
 end
 for B in (
diff --git a/src/Types.jl b/src/Types.jl
index 4ed61f2f19..c909e5395e 100644
--- a/src/Types.jl
+++ b/src/Types.jl
@@ -108,6 +108,9 @@ mutable struct TracedRArray{T,N,RT<:TracedRNumber{T}} <: RArray{RT,N}
     function TracedRArray{T,N,RT}(
         paths::Tuple, mlir_data::Union{Nothing,MLIR.IR.Value}, shape
     ) where {T,N,RT<:TracedRNumber{T}}
+        RT === traced_number_type(T) ||
+            throw(ArgumentError("Element type of TracedRArray{$T} must be \
+                                 $(traced_number_type(T)), got $RT"))
         shape = Tuple(shape)
         if !isnothing(mlir_data)
             @assert size(MLIR.IR.type(mlir_data)) == shape "Expected: $(shape), got: $(size(MLIR.IR.type(mlir_data)))"
@@ -130,6 +133,18 @@ function TracedRArray{T,N}(::UndefInitializer, shape::NTuple{N,Int}) where {T,N}
     return similar(TracedRArray{T,N}, shape)
 end
 
+function TracedRArray{T,N,RT}(
+    ::UndefInitializer, shape::Integer...
+) where {T,N,RT<:TracedRNumber{T}}
+    return similar(TracedRArray{T,N}, shape...)
+end
+
+function TracedRArray{T,N,RT}(
+    ::UndefInitializer, shape::NTuple{N,Int}
+) where {T,N,RT<:TracedRNumber{T}}
+    return similar(TracedRArray{T,N}, shape)
+end
+
 function repath(x::TracedRArray{T,N}, paths) where {T,N}
     return Core.Typeof(x)(paths, x.mlir_data, x.shape)
 end
@@ -274,7 +289,7 @@ end
 # `convert(ConcretePJRTFloat{Float64,1}, x)`. The `Rational`/`Complex` methods
 # disambiguate against Base's constructors on `Integer`/`AbstractFloat`/`Real`.
 for CT in (:ConcretePJRTInteger, :ConcretePJRTFloat, :ConcretePJRTComplex),
-    XT in (:Number, :Rational, :Complex)
+    XT in (:Number, :Rational, :Complex, :BigFloat)
 
     @eval function (::Type{CT})(x::$XT; kwargs...) where {CT<:$CT}
         return ConcretePJRTNumber{unwrapped_eltype(CT)}(x; kwargs...)
@@ -443,7 +458,7 @@ function ConcreteIFRTNumber(data::T; kwargs...) where {T<:Number}
 end
 
 for CT in (:ConcreteIFRTInteger, :ConcreteIFRTFloat, :ConcreteIFRTComplex),
-    XT in (:Number, :Rational, :Complex)
+    XT in (:Number, :Rational, :Complex, :BigFloat)
 
     @eval function (::Type{CT})(x::$XT; kwargs...) where {CT<:$CT}
         return ConcreteIFRTNumber{unwrapped_eltype(CT)}(x; kwargs...)
diff --git a/test/integration/structarrays.jl b/test/integration/structarrays.jl
index 4ff188f166..69eb884be3 100644
--- a/test/integration/structarrays.jl
+++ b/test/integration/structarrays.jl
@@ -9,7 +9,7 @@ using StructArrays, StaticArrays, Reactant, LinearAlgebra, Test
     # Note that the element type (the NamedTuple) contains ConcreteRNumbers even though track_numbers is not enabled.
     # This is because when the backing arrays are converted to TracedRArrays, their elements will contain TracedRNumbers.
     # In order for the element type to match the backing arrays, we need to use ConcreteRNumbers here as well:
-    @test typeof(x_ra) <: (StructArray{
+    @test typeof(x_ra)<:(StructArray{
         <:NamedTuple{
             (:a, :b, :c),
             <:Tuple{<:ConcreteRNumber{Float64},String,<:ConcreteRNumber{Float32}},
@@ -26,15 +26,13 @@ using StructArrays, StaticArrays, Reactant, LinearAlgebra, Test
         Reactant.make_tracer(Reactant.OrderedIdDict(), x_ra, (), Reactant.ConcreteToTraced)
     ) == StructArray{
         @NamedTuple{
-            a::Reactant.TracedRNumber{Float64},
-            b::String,
-            c::Reactant.TracedRNumber{Float32},
+            a::Reactant.TracedRFloat{Float64},b::String,c::Reactant.TracedRFloat{Float32}
         },
         2,
         @NamedTuple{
-            a::Reactant.TracedRArray{Float64,2},
+            a::Reactant.TracedRArray{Float64,2,Reactant.TracedRFloat{Float64}},
             b::Matrix{String},
-            c::Reactant.TracedRArray{Float32,2},
+            c::Reactant.TracedRArray{Float32,2,Reactant.TracedRFloat{Float32}},
         },
         Int64,
     }

From a6ff657eef9fb815055675aa59560a5c6305370f Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Wed, 8 Jul 2026 11:53:01 +0200
Subject: [PATCH 15/30] Adopt Enzyme's upstream wrapped-number activity
 classification

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.
---
 src/Enzyme.jl         | 18 +-----------------
 test/core/autodiff.jl | 12 ++++++++++++
 2 files changed, 13 insertions(+), 17 deletions(-)

diff --git a/src/Enzyme.jl b/src/Enzyme.jl
index 60162295f5..2dfba1706b 100644
--- a/src/Enzyme.jl
+++ b/src/Enzyme.jl
@@ -67,22 +67,6 @@ for RT in (:RInteger, :RFloat, :RComplex)
     @eval @inline Enzyme.make_zero(x::$RT) = zero(Core.Typeof(x))
 end
 
-# Reactant numbers are mutable wrappers around device/MLIR values and must not
-# be classified as active scalars, which Enzyme does for any `AbstractFloat`.
-function Enzyme.Compiler.active_reg_nothrow(::Type{T}) where {T<:RNumber}
-    return Enzyme.Compiler.DupState
-end
-
-# Reactant numbers are mutable wrappers around device/MLIR values and must
-# stay `Duplicated` under autodiff even though they subtype `AbstractFloat`.
-for RT in (:RArray, :RNumber, :RInteger, :RFloat, :RComplex)
-    @eval @inline function Enzyme.guess_activity(
-        ::Type{T}, mode::Enzyme.Mode
-    ) where {T<:$RT}
-        return Enzyme.Duplicated{T}
-    end
-end
-
 @inline function Enzyme.make_zero(x::RArray{FT,N})::RArray{FT,N} where {FT<:AbstractFloat,N}
     return Base.zero(x)
 end
@@ -603,7 +587,7 @@ function overload_autodiff(
         if arg isa Active && width == 1 && isempty(path[3:end])
             # `Active` arguments return their gradient in the result tuple
             # instead of accumulating into a shadow
-            restup[idx - fnwrap] = Reactant.TracedRNumber{unwrapped_eltype(arg.val)}(
+            restup[idx - fnwrap] = TracedRNumber{unwrapped_eltype(arg.val)}(
                 (), TracedUtils.transpose_val(MLIR.IR.result(res, residx))
             )
         else
diff --git a/test/core/autodiff.jl b/test/core/autodiff.jl
index f0779be5a1..514ae9f4af 100644
--- a/test/core/autodiff.jl
+++ b/test/core/autodiff.jl
@@ -173,6 +173,18 @@ end
     @test res2 ≈ 4 * 3 * 3.1^2
 end
 
+@testset "Explicit Active scalar arguments" begin
+    x = ConcreteRNumber(3.1)
+    f(x) = x * x * x * x
+    df(x) = Enzyme.autodiff(Reverse, f, Active, Active(x))[1][1]
+    res1 = @jit df(x)
+    @test res1 ≈ 4 * 3.1^3
+    dfp(x) = Enzyme.autodiff(ReverseWithPrimal, f, Active, Active(x))
+    res2 = @jit dfp(x)
+    @test res2[1][1] ≈ 4 * 3.1^3
+    @test res2[2] ≈ 3.1^4
+end
+
 @testset "Seed initialization of Complex arrays on matmul: Issue #593" begin
     df(x, y) = Enzyme.gradient(ReverseWithPrimal, *, x, y)
     @test begin

From 4c17102542d13f0e75a05c8c0009574e57cfa86a Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Wed, 8 Jul 2026 11:53:29 +0200
Subject: [PATCH 16/30] Resolve method ambiguities and traced-rational
 promotion gaps

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.
---
 .../ReactantSparseArraysExt.jl                | 33 +++++++++-
 ext/ReactantSparseArraysExt/ReadOnly.jl       |  7 ++
 src/ConcreteRArray.jl                         |  2 +-
 src/Indexing.jl                               | 27 ++++----
 src/TracedRArray.jl                           | 15 ++++-
 src/TracedRNumber.jl                          | 65 +++++++++++++++----
 src/TracedRational.jl                         | 18 +++++
 src/Types.jl                                  | 26 +++++++-
 src/compiler/Codegen.jl                       | 16 ++---
 9 files changed, 169 insertions(+), 40 deletions(-)

diff --git a/ext/ReactantSparseArraysExt/ReactantSparseArraysExt.jl b/ext/ReactantSparseArraysExt/ReactantSparseArraysExt.jl
index b782ad62ad..8c2d1c1589 100644
--- a/ext/ReactantSparseArraysExt/ReactantSparseArraysExt.jl
+++ b/ext/ReactantSparseArraysExt/ReactantSparseArraysExt.jl
@@ -1,10 +1,39 @@
 module ReactantSparseArraysExt
 
-using Reactant: Reactant, TracedRNumber
+using Reactant: Reactant, TracedRInteger, TracedRNumber
 using SparseArrays:
-    SparseArrays, ReadOnly, AbstractSparseArray, CHOLMOD, AbstractSparseMatrixCSC
+    SparseArrays,
+    ReadOnly,
+    AbstractSparseArray,
+    AbstractSparseVector,
+    CHOLMOD,
+    AbstractSparseMatrixCSC
 
 include("Errors.jl")
 include("ReadOnly.jl")
 
+# Scalar indexing disambiguators against the generic traced-array `getindex`;
+# they delegate to the SparseArrays implementations.
+function Base.getindex(
+    x::AbstractSparseMatrixCSC{<:TracedRNumber},
+    i::Union{Int,TracedRInteger{Int}},
+    j::Union{Int,TracedRInteger{Int}},
+)
+    return Base.invoke(getindex, Tuple{AbstractSparseMatrixCSC,Integer,Integer}, x, i, j)
+end
+
+function Base.getindex(
+    x::AbstractSparseVector{<:TracedRNumber}, i::Union{Int,TracedRInteger{Int}}
+)
+    return Base.invoke(getindex, Tuple{AbstractSparseVector,Integer}, x, i)
+end
+
+function Base.getindex(
+    x::CHOLMOD.Sparse{<:TracedRNumber},
+    i::Union{Int,TracedRInteger{Int}},
+    j::Union{Int,TracedRInteger{Int}},
+)
+    return Base.invoke(getindex, Tuple{CHOLMOD.Sparse,Integer,Integer}, x, i, j)
+end
+
 end
diff --git a/ext/ReactantSparseArraysExt/ReadOnly.jl b/ext/ReactantSparseArraysExt/ReadOnly.jl
index 3903a405fb..18685535cd 100644
--- a/ext/ReactantSparseArraysExt/ReadOnly.jl
+++ b/ext/ReactantSparseArraysExt/ReadOnly.jl
@@ -48,6 +48,13 @@ function Base.getindex(
     return getindex(parent(x), idx...)
 end
 
+# Disambiguates against the generic traced-array `getindex`.
+function Base.getindex(
+    x::ReadOnly{<:TracedRNumber{T},N}, idx::Vararg{Union{Int,Reactant.TracedRNumber{Int}},N}
+) where {T,N}
+    return getindex(parent(x), idx...)
+end
+
 function Base.getindex(
     x::SparseArrays.ReadOnly{<:Reactant.TracedRNumber{T},1,V},
     idx::Union{Int64,Reactant.TracedRNumber{Int64}},
diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 4b444abc7a..b94d97a426 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -200,7 +200,7 @@ end
 
 function Base.convert(
     ::Type{T}, x::Union{ConcretePJRTScalar{T},ConcreteIFRTScalar{T}}
-) where {T<:Number}
+) where {T<:ReactantPrimitive}
     return to_number(x)
 end
 function Base.convert(
diff --git a/src/Indexing.jl b/src/Indexing.jl
index 0461f8814f..d2a29c3347 100644
--- a/src/Indexing.jl
+++ b/src/Indexing.jl
@@ -109,6 +109,7 @@ for (aT, iT) in (
     (Base.OneTo, TracedRInteger),
     (Base.OneTo{<:TracedRNumber}, Int),
     (Base.OneTo{<:TracedRNumber}, TracedRInteger),
+    (Base.OneTo{<:TracedRNumber}, TracedRInteger{Int}),
 )
     @eval function Base.getindex(a::$aT, index::$iT)
         return convert(TracedRNumber{Reactant.unwrapped_eltype(a)}, index)
@@ -117,8 +118,8 @@ end
 
 for (aT, iT) in (
     (StepRangeLen, TracedRInteger),
-    (StepRangeLen{<:TracedRNumber}, Int),
     (StepRangeLen{<:TracedRNumber}, TracedRInteger),
+    (StepRangeLen{<:TracedRNumber}, TracedRInteger{Int}),
 )
     @eval function Base.getindex(r::$aT, index::$iT)
         # FIXME: this crashes for some reason
@@ -132,8 +133,8 @@ end
 
 for (aT, iT) in (
     (LinRange, TracedRInteger),
-    (LinRange{<:TracedRNumber}, Int),
     (LinRange{<:TracedRNumber}, TracedRInteger),
+    (LinRange{<:TracedRNumber}, TracedRInteger{Int}),
 )
     @eval function Base.getindex(r::$aT, index::$iT)
         return convert(
@@ -150,16 +151,6 @@ function Base.getindex(
     return getindex(Reactant.promote_to(TracedRArray{T,1}, a), index)
 end
 
-function Base.getindex(
-    a::Union{
-        LinRange{T} where T<:Reactant.TracedRNumber,
-        StepRangeLen{T} where T<:Reactant.TracedRNumber,
-    },
-    index::Int,
-)
-    return getindex(Reactant.promote_to(TracedRArray{T,1}, a), index)
-end
-
 function Base.getindex(a::AnyTracedRArray{T,N}, linear_indices) where {T,N}
     ancestor, idxs = TracedUtils.get_ancestor_and_indices(a, linear_indices)
     return getindex(ancestor, idxs...)
@@ -207,7 +198,7 @@ function Base.getindex(
 end
 
 function Base.getindex(
-    x::Base.Sort.WithoutMissingVector{<:TracedRNumber{T}}, i::Int
+    x::Base.Sort.WithoutMissingVector{<:TracedRNumber{T}}, i::Union{Int,TracedRInteger{Int}}
 ) where {T}
     out = getindex(x.data, i)
     @assert !(out isa Missing)
@@ -221,6 +212,14 @@ end
 function Base.getindex(v::TracedUnitRange{<:TracedRNumber}, i::TracedRInteger{Int})
     return convert(eltype(v), v.start + (i - oneunit(i)))
 end
+function Base.getindex(
+    v::TracedUnitRange{<:TracedRNumber}, i::Union{Int,TracedRInteger{Int}}
+)
+    return convert(eltype(v), v.start + (i - oneunit(i)))
+end
+function Base.getindex(v::TracedUnitRange{<:TracedRNumber}, i::Int)
+    return convert(eltype(v), v.start + (i - oneunit(i)))
+end
 function Base.getindex(r::TracedStepRangeLen{<:TracedRNumber}, i::TracedRInteger{Int})
     return Base.unsafe_getindex(r, i)
 end
@@ -328,7 +327,7 @@ end
 
 Base.getindex(v::TracedUnitRange, ::Colon) = v
 
-for iType in (Int, TracedRNumber{Int}, Integer)
+for iType in (Int, TracedRNumber{Int}, Integer, TracedRInteger, TracedRInteger{Int})
     @eval function Base.getindex(v::TracedUnitRange{T}, i::$iType) where {T}
         return convert(T, v.start + (i - oneunit(i)))
         # TODO(#2237): we should have error messages at some point.
diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
index 29d2379f12..7517984f57 100644
--- a/src/TracedRArray.jl
+++ b/src/TracedRArray.jl
@@ -4,7 +4,13 @@ using Base: Broadcast
 using Base.Broadcast: Broadcasted, AbstractArrayStyle, instantiate
 
 using ..Reactant:
-    Reactant, TracedRArray, TracedRNumber, TracedRInteger, AnyTracedRArray, AnyTracedRVector
+    Reactant,
+    TracedRArray,
+    TracedRNumber,
+    TracedRInteger,
+    TracedRReal,
+    AnyTracedRArray,
+    AnyTracedRVector
 using ..Reactant: MLIR, unwrapped_eltype
 using ..Ops: @opcall
 using ..TracedUtils: TracedUtils, get_mlir_data, set_mlir_data!, materialize_traced_array
@@ -1111,13 +1117,18 @@ end
 
 for (aType, xType) in (
     (AbstractRange{<:Real}, TracedRNumber{<:Real}),
+    (AbstractRange{<:Integer}, TracedRReal{<:Real}),
+    (AbstractRange{<:TracedRInteger}, TracedRReal{<:Real}),
     (AbstractRange{<:TracedRNumber}, Real),
     (AbstractRange{<:TracedRNumber}, TracedRNumber{<:Real}),
 )
     @eval function Base.searchsortedfirst(
         a::$(aType), x::$(xType), o::Base.DirectOrdering
     )::TracedRNumber{keytype(a)}
-        x = TracedUtils.promote_to(TracedRNumber{Reactant.unwrapped_eltype(a)}, x)
+        # promoting `x` to the range element type would truncate, e.g. for
+        # `searchsortedfirst(1:10, 2.5)`
+        T = Base.promote_type(Reactant.unwrapped_eltype(a), Reactant.unwrapped_eltype(x))
+        x = TracedUtils.promote_to(TracedRNumber{T}, x)
 
         f, h, l = first(a), step(a), last(a)
         n = round(Int, (x - f) / h + 1)
diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index 7b0677ea76..f59570e5c5 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -150,17 +150,45 @@ for S in Base.uniontypes(Reactant.ReactantFloat8)
     end
 end
 
-# Disambiguate against Base's `Rational` and `Complex` promotion rules.
+# Base defines its own `BigInt`/`BigFloat` promotion rules per numeric kind;
+# these methods only disambiguate against them.
+function Base.promote_rule(::Type{BigInt}, ::Type{<:TracedRInteger{T}}) where {T}
+    return traced_number_type(Base.promote_type(BigInt, T))
+end
+function Base.promote_rule(::Type{BigInt}, ::Type{<:TracedRFloat{T}}) where {T}
+    return traced_number_type(Base.promote_type(BigInt, T))
+end
+function Base.promote_rule(::Type{BigFloat}, ::Type{<:TracedRReal{T}}) where {T}
+    return traced_number_type(Base.promote_type(BigFloat, T))
+end
+function Base.promote_rule(::Type{BigFloat}, ::Type{<:TracedRFloat{T}}) where {T}
+    return traced_number_type(Base.promote_type(BigFloat, T))
+end
+
+# Disambiguate against Base's `Rational` and `Complex` promotion rules. Both
+# argument orders are needed: the generic rules above intercept the reverse
+# direction before Base's `Bottom` fallback, and their result computation
+# cannot represent rationals.
 function Base.promote_rule(
     ::Type{Rational{T}}, ::Type{<:TracedRInteger{S}}
 ) where {T<:Integer,S}
     return Reactant.TracedRational{traced_number_type(Base.promote_type(T, S))}
 end
+function Base.promote_rule(
+    ::Type{<:TracedRInteger{S}}, ::Type{Rational{T}}
+) where {T<:Integer,S}
+    return Reactant.TracedRational{traced_number_type(Base.promote_type(T, S))}
+end
 function Base.promote_rule(
     ::Type{Rational{T}}, ::Type{<:TracedRFloat{S}}
 ) where {T<:Integer,S}
     return traced_number_type(Base.promote_type(Rational{T}, S))
 end
+function Base.promote_rule(
+    ::Type{<:TracedRFloat{S}}, ::Type{Rational{T}}
+) where {T<:Integer,S}
+    return traced_number_type(Base.promote_type(Rational{T}, S))
+end
 function Base.promote_rule(::Type{Complex{T}}, ::Type{<:TracedRReal{S}}) where {T<:Real,S}
     return traced_number_type(Base.promote_type(Complex{T}, S))
 end
@@ -254,10 +282,14 @@ for (RT, XT) in (
     (:TracedRInteger, :Complex),
     (:TracedRFloat, :Complex),
     (:TracedRInteger, :Rational),
-    (:TracedRFloat, :Rational),
+    (:TracedRInteger, :BigFloat),
 )
     @eval (::Type{RT})(x::$XT) where {RT<:$RT} = Reactant.promote_to(RT, x)
 end
+# The ambiguity resolver requires the exact TypeVar shape of the competing
+# Base constructor: `(::Type{T})(::Rational{S}) where {S,T<:AbstractFloat}`
+# is `S`-parameterized, `(::Type{T})(::Rational) where {T<:Integer}` is not.
+(::Type{RT})(x::Rational{S}) where {S,RT<:TracedRFloat} = Reactant.promote_to(RT, x)
 
 for T in Base.uniontypes(Reactant.ReactantFloat8)
     @eval TracedRNumber{T}(x::$T) where {T} = Reactant.promote_to(TracedRNumber{T}, x)
@@ -505,12 +537,13 @@ function Base.div(
         signbit(lhs) == signbit(rhs), div(lhs, rhs, RoundUp), div(lhs, rhs, RoundDown)
     )
 end
-function Base.div(
-    @nospecialize(lhs::TracedRInteger{T}),
-    @nospecialize(rhs::TracedRInteger{T}),
-    r::Union{typeof(RoundNearest),typeof(RoundNearestTiesAway),typeof(RoundNearestTiesUp)},
-) where {T<:Integer}
-    return divrem(lhs, rhs, r)[1]
+for RM in
+    (RoundingMode{:Nearest}, RoundingMode{:NearestTiesAway}, RoundingMode{:NearestTiesUp})
+    @eval function Base.div(
+        @nospecialize(lhs::TracedRInteger{T}), @nospecialize(rhs::TracedRInteger{T}), r::$RM
+    ) where {T<:Integer}
+        return divrem(lhs, rhs, r)[1]
+    end
 end
 
 function Base.div(
@@ -858,9 +891,12 @@ Base.:!(x::TracedRInteger) = @opcall not(x)
 # With a traced integer exponent, Base's `^(::Number, ::Integer)` would call
 # `power_by_squaring`, which branches on traced booleans. The extra methods
 # disambiguate against Base's specialized `^` methods.
-for B in (:TracedRInteger, :TracedRFloat, :TracedRComplex, :Real, :Complex, :Rational)
+for B in (:TracedRInteger, :TracedRFloat, :TracedRComplex, :Real, :Complex)
     @eval Base.:^(x::$B, p::TracedRInteger) = ^(promote(x, p)...)
 end
+# A traced exponent's sign is unknown at trace time, so rationality cannot be
+# preserved (`x^p` vs `inv(x)^-p`); compute in floating point instead.
+Base.:^(x::Rational, p::TracedRInteger) = ^(float(x), p)
 for B in (
     Float16,
     Float32,
@@ -870,10 +906,19 @@ for B in (
     Complex{<:AbstractFloat},
     Complex{<:Integer},
     Complex{<:Rational},
+    Irrational{:ℯ},
+    BigInt,
+    BigFloat,
 )
     @eval Base.:^(x::$B, p::TracedRInteger) = ^(promote(x, p)...)
 end
 
+# Same-`T` disambiguators against the same-type traced `^` method.
+Base.:^(x::TracedRInteger{T}, p::TracedRInteger{T}) where {T} = @opcall power(x, p)
+Base.:^(x::TracedRFloat{T}, p::TracedRInteger{T}) where {T} = ^(promote(x, p)...)
+Base.:^(x::TracedRComplex{T}, p::TracedRInteger{T}) where {T} = ^(promote(x, p)...)
+Base.:^(x::TracedRReal{T}, p::TracedRInteger{T}) where {T} = ^(promote(x, p)...)
+
 Base.fma(x::TracedRFloat{T}, y::TracedRFloat{T}, z::TracedRFloat{T}) where {T} = x * y + z
 
 function Base.literal_pow(
@@ -1085,8 +1130,6 @@ for TracedT in (:TracedRInteger, :TracedRFloat, :TracedRComplex)
     @eval Base.float(::Type{<:$TracedT{T}}) where {T} = traced_number_type(float(T))
 end
 
-using Reactant: ReactantFloat, ReactantInt
-
 Base.round(A::TracedRFloat) = @opcall round_nearest_even(A)
 Base.round(A::TracedRInteger) = A
 function Base.round(A::TracedRFloat, ::typeof(RoundNearest))
diff --git a/src/TracedRational.jl b/src/TracedRational.jl
index d1cfc59734..08fd5f0e32 100644
--- a/src/TracedRational.jl
+++ b/src/TracedRational.jl
@@ -84,6 +84,11 @@ end
 function Base.promote_rule(::Type{TracedRational{T}}, ::Type{Rational{S}}) where {T,S}
     return TracedRational{promote_type(T, S)}
 end
+function Base.promote_rule(
+    ::Type{TracedRational{T}}, ::Type{<:Reactant.TracedRInteger{S}}
+) where {T,S}
+    return TracedRational{promote_type(T, Reactant.TracedRInteger{S})}
+end
 function Base.promote_rule(::Type{TracedRational{T}}, ::Type{S}) where {T,S<:AbstractFloat}
     return S
 end
@@ -113,6 +118,9 @@ end
 function Base.copysign(x::TracedRational, y::Rational)
     return TracedRational(copysign(x.num, y.num), x.den)
 end
+function Base.copysign(x::TracedRational, y::Reactant.TracedRReal)
+    return TracedRational(copysign(x.num, y), x.den)
+end
 
 Base.abs(x::TracedRational) = TracedRational(abs(x.num), x.den)
 
@@ -151,6 +159,16 @@ for op in (:+, :-, :rem, :mod)
     end
 end
 
+# `TracedRInteger <: Integer`, so Base's checked `(Integer, Rational)`
+# arithmetic would apply; promote to `TracedRational` instead. (`mod` and
+# `rem` for these pairs are defined in TracedRNumber.jl.)
+for op in (:+, :-, :*)
+    @eval begin
+        Base.$(op)(x::Reactant.TracedRInteger, y::Rational) = $(op)(promote(x, y)...)
+        Base.$(op)(x::Rational, y::Reactant.TracedRInteger) = $(op)(promote(x, y)...)
+    end
+end
+
 function Base.:*(x::TracedRational, y::TracedRational)
     xn, yd = Base.divgcd(promote(x.num, y.den)...)
     yn, xd = Base.divgcd(promote(y.num, x.den)...)
diff --git a/src/Types.jl b/src/Types.jl
index c909e5395e..30627cb52c 100644
--- a/src/Types.jl
+++ b/src/Types.jl
@@ -289,12 +289,23 @@ end
 # `convert(ConcretePJRTFloat{Float64,1}, x)`. The `Rational`/`Complex` methods
 # disambiguate against Base's constructors on `Integer`/`AbstractFloat`/`Real`.
 for CT in (:ConcretePJRTInteger, :ConcretePJRTFloat, :ConcretePJRTComplex),
-    XT in (:Number, :Rational, :Complex, :BigFloat)
+    XT in (:Number, :Complex, :BigFloat)
 
     @eval function (::Type{CT})(x::$XT; kwargs...) where {CT<:$CT}
         return ConcretePJRTNumber{unwrapped_eltype(CT)}(x; kwargs...)
     end
 end
+for CT in (:ConcretePJRTInteger, :ConcretePJRTComplex)
+    @eval function (::Type{CT})(x::Rational; kwargs...) where {CT<:$CT}
+        return ConcretePJRTNumber{unwrapped_eltype(CT)}(x; kwargs...)
+    end
+end
+# The ambiguity resolver requires the exact TypeVar shape of the competing
+# Base constructor: `(::Type{T})(::Rational{S}) where {S,T<:AbstractFloat}`
+# is `S`-parameterized, `(::Type{T})(::Rational) where {T<:Integer}` is not.
+function (::Type{CT})(x::Rational{S}; kwargs...) where {S,CT<:ConcretePJRTFloat}
+    return ConcretePJRTNumber{unwrapped_eltype(CT)}(x; kwargs...)
+end
 
 function ConcretePJRTNumber(data::ConcretePJRTNumber; kwargs...)
     return ConcretePJRTNumber(
@@ -458,12 +469,23 @@ function ConcreteIFRTNumber(data::T; kwargs...) where {T<:Number}
 end
 
 for CT in (:ConcreteIFRTInteger, :ConcreteIFRTFloat, :ConcreteIFRTComplex),
-    XT in (:Number, :Rational, :Complex, :BigFloat)
+    XT in (:Number, :Complex, :BigFloat)
 
     @eval function (::Type{CT})(x::$XT; kwargs...) where {CT<:$CT}
         return ConcreteIFRTNumber{unwrapped_eltype(CT)}(x; kwargs...)
     end
 end
+for CT in (:ConcreteIFRTInteger, :ConcreteIFRTComplex)
+    @eval function (::Type{CT})(x::Rational; kwargs...) where {CT<:$CT}
+        return ConcreteIFRTNumber{unwrapped_eltype(CT)}(x; kwargs...)
+    end
+end
+# The ambiguity resolver requires the exact TypeVar shape of the competing
+# Base constructor: `(::Type{T})(::Rational{S}) where {S,T<:AbstractFloat}`
+# is `S`-parameterized, `(::Type{T})(::Rational) where {T<:Integer}` is not.
+function (::Type{CT})(x::Rational{S}; kwargs...) where {S,CT<:ConcreteIFRTFloat}
+    return ConcreteIFRTNumber{unwrapped_eltype(CT)}(x; kwargs...)
+end
 
 function ConcreteIFRTNumber(data::ConcreteIFRTNumber; kwargs...)
     return ConcreteIFRTNumber(
diff --git a/src/compiler/Codegen.jl b/src/compiler/Codegen.jl
index 6213e851f1..996df09d32 100644
--- a/src/compiler/Codegen.jl
+++ b/src/compiler/Codegen.jl
@@ -366,14 +366,14 @@ end
 # method) so that integer/float scalars are not ambiguous with the
 # `create_result(::Union{Integer,AbstractFloat,...})` method below.
 for (NumT, worker) in (
-    (:ConcretePJRTNumber, :create_pjrt_number_result),
-    (:ConcretePJRTInteger, :create_pjrt_number_result),
-    (:ConcretePJRTFloat, :create_pjrt_number_result),
-    (:ConcretePJRTComplex, :create_pjrt_number_result),
-    (:ConcreteIFRTNumber, :create_ifrt_number_result),
-    (:ConcreteIFRTInteger, :create_ifrt_number_result),
-    (:ConcreteIFRTFloat, :create_ifrt_number_result),
-    (:ConcreteIFRTComplex, :create_ifrt_number_result),
+    (ConcretePJRTNumber, create_pjrt_number_result),
+    (ConcretePJRTInteger, create_pjrt_number_result),
+    (ConcretePJRTFloat, create_pjrt_number_result),
+    (ConcretePJRTComplex, create_pjrt_number_result),
+    (ConcreteIFRTNumber, create_ifrt_number_result),
+    (ConcreteIFRTInteger, create_ifrt_number_result),
+    (ConcreteIFRTFloat, create_ifrt_number_result),
+    (ConcreteIFRTComplex, create_ifrt_number_result),
 )
     @eval function create_result(
         tocopy::$NumT,

From 78cb028018c62be69bd6a267c5c03f04b106398f Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Wed, 8 Jul 2026 11:53:37 +0200
Subject: [PATCH 17/30] Document the number kind hierarchy

Created by generative AI.
---
 docs/src/api/api.md |  8 ++++++++
 src/Types.jl        | 37 ++++++++++++++++++++++++++++++++++---
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/docs/src/api/api.md b/docs/src/api/api.md
index 2e25512a57..9674527197 100644
--- a/docs/src/api/api.md
+++ b/docs/src/api/api.md
@@ -34,6 +34,14 @@ Reactant.to_rarray
 ```@docs
 ConcreteRArray
 ConcreteRNumber
+Reactant.RNumber
+Reactant.RReal
+RInteger
+RFloat
+RComplex
+Reactant.traced_number_type
+Reactant.pjrt_number_type
+Reactant.ifrt_number_type
 ```
 
 ## Inspect Generated HLO
diff --git a/src/Types.jl b/src/Types.jl
index 30627cb52c..bf2b88e1a4 100644
--- a/src/Types.jl
+++ b/src/Types.jl
@@ -1,11 +1,42 @@
-# Reactant numbers are split by numeric kind so that traced and concrete values
-# subtype `Integer` and `AbstractFloat` where possible. Julia has no abstract
-# complex type, so complex Reactant numbers can only subtype `Number`.
+"""
+    RInteger{T} <: Integer
+
+Abstract supertype of Reactant integers (traced and concrete) with underlying
+element type `T`.
+"""
 abstract type RInteger{T<:ReactantPrimitive} <: Integer end
+
+"""
+    RFloat{T} <: AbstractFloat
+
+Abstract supertype of Reactant floats (traced and concrete) with underlying
+element type `T`.
+"""
 abstract type RFloat{T<:ReactantPrimitive} <: AbstractFloat end
+
+"""
+    RComplex{T} <: Number
+
+Abstract supertype of complex Reactant numbers (traced and concrete) with
+underlying element type `T`. Julia has no abstract complex type, so complex
+Reactant numbers can only subtype `Number`.
+"""
 abstract type RComplex{T<:ReactantPrimitive} <: Number end
 
+"""
+    RReal{T}
+
+Union of the real Reactant number kinds `RInteger{T}` and `RFloat{T}`.
+"""
 const RReal{T} = Union{RInteger{T},RFloat{T}}
+
+"""
+    RNumber{T}
+
+Union of all Reactant number kinds: `RInteger{T}`, `RFloat{T}` and
+`RComplex{T}`. Dispatch on the kind types to target a subset of Reactant
+numbers.
+"""
 const RNumber{T} = Union{RReal{T},RComplex{T}}
 
 abstract type AbstractConcreteInteger{T} <: RInteger{T} end

From 243e5ad95f84b3eb49eb6a1bb3d2087fbe69fae0 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Wed, 8 Jul 2026 16:01:15 +0200
Subject: [PATCH 18/30] Format

Created by generative AI.
---
 ext/ReactantCUDAExt.jl           |   6 +-
 src/ConcreteRArray.jl            |   2 +-
 src/Indexing.jl                  |   4 +-
 src/Ops.jl                       | 173 ++++++++++++++++---------------
 src/TracedRArray.jl              |   2 +-
 src/TracedUtils.jl               |  19 ++--
 src/Types.jl                     |  12 ++-
 src/compiler/Codegen.jl          |  99 ++++++++++++------
 test/core/compile.jl             |   9 +-
 test/integration/structarrays.jl |   4 +-
 10 files changed, 188 insertions(+), 142 deletions(-)

diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl
index 603860732f..5aae5f9d37 100644
--- a/ext/ReactantCUDAExt.jl
+++ b/ext/ReactantCUDAExt.jl
@@ -864,9 +864,7 @@ function compile(job)
             prevdl = String(prevdlattr)
             @assert prevdl == dl "data layout mismatch, tried compiling cuda kernels for different target machines?"
         else
-            MLIR.IR.setattr!(
-                MLIR.IR.Operation(cur_module), dl_attr_name, MLIR.IR.Attribute(dl)
-            )
+            MLIR.IR.setattr!(MLIR.IR.Operation(cur_module), dl_attr_name, MLIR.IR.Attribute(dl))
         end
 
         return String(Reactant.TracedUtils.get_attribute_by_name(linkRes, "sym_name"))
@@ -1119,7 +1117,7 @@ Reactant.@reactant_overlay function (func::LLVMFunc{F,tt})(
     has_cast_float_type =
         bfloat16_compile_type !== BFloat16 && any(values(seen)) do arg
             (arg isa TracedRArray || arg isa TracedRNumber) &&
-            Reactant.unwrapped_eltype(typeof(arg)) === BFloat16
+                Reactant.unwrapped_eltype(typeof(arg)) === BFloat16
         end
 
     wrapper_tys = MLIR.IR.Type[]
diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index b94d97a426..eef3db099e 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -665,7 +665,7 @@ function Base.copy(bc::Base.Broadcast.Broadcasted{Broadcast.ArrayStyle{ConcreteP
     if all(buffer_on_cpu, bc.args) && all(
         x ->
             !(x isa ConcretePJRTArray) ||
-            (x isa ConcretePJRTArray && !Sharding.is_sharded(x)),
+                (x isa ConcretePJRTArray && !Sharding.is_sharded(x)),
         bc.args,
     )
         ElType = Base.Broadcast.combine_eltypes(bc.f, bc.args)
diff --git a/src/Indexing.jl b/src/Indexing.jl
index d2a29c3347..2b71a60ac2 100644
--- a/src/Indexing.jl
+++ b/src/Indexing.jl
@@ -458,7 +458,9 @@ function Base.setindex!(a::TracedRArray{T,N}, v, indices::Vararg{Any,N}) where {
 
     TracedUtils.set_mlir_data!(
         a,
-        @opcall(dynamic_update_slice(a, v, [i isa Colon ? 1 : first(i) for i in indices])).mlir_data,
+        @opcall(
+            dynamic_update_slice(a, v, [i isa Colon ? 1 : first(i) for i in indices])
+        ).mlir_data,
     )
     return v
 end
diff --git a/src/Ops.jl b/src/Ops.jl
index 807df960ff..0f8841f9ec 100644
--- a/src/Ops.jl
+++ b/src/Ops.jl
@@ -45,10 +45,12 @@ function _function_macro_error()
 end
 
 macro caller_function()
-    return esc(quote
-        $(Expr(:isdefined, :var"#self#")) || $(_function_macro_error)()
-        var"#self#"
-    end)
+    return esc(
+        quote
+            $(Expr(:isdefined, :var"#self#")) || $(_function_macro_error)()
+            var"#self#"
+        end,
+    )
 end
 
 """
@@ -1377,15 +1379,16 @@ end
         sample_inputs[2i - 1] = Reactant.promote_to(TracedRNumber{T}, 0)
         sample_inputs[2i] = Reactant.promote_to(TracedRNumber{T}, 0)
     end
-    func = Reactant.TracedUtils.make_mlir_fn(
-        comparator,
-        (sample_inputs...,),
-        (),
-        "comparator",
-        false;
-        args_in_result=:none,
-        return_dialect=:stablehlo,
-    ).f
+    func =
+        Reactant.TracedUtils.make_mlir_fn(
+            comparator,
+            (sample_inputs...,),
+            (),
+            "comparator",
+            false;
+            args_in_result=:none,
+            return_dialect=:stablehlo,
+        ).f
     @assert MLIR.IR.nregions(func) == 1
     fn_name = String(
         MLIR.IR.getattr(func, String(MLIR.API.mlirSymbolTableGetSymbolAttributeName()))
@@ -1434,20 +1437,21 @@ end
 ) where {T<:AbstractFloat,N}
     fallback === missing && (fallback = Reactant.FALLBACK_APPROX_TOP_K_LOWERING[])
 
-    func = Reactant.TracedUtils.make_mlir_fn(
-        comparator,
-        (
-            Reactant.promote_to(TracedRNumber{T}, 0),
-            Reactant.promote_to(TracedRNumber{T}, 0),
-            Reactant.promote_to(TracedRNumber{Int32}, 0),
-            Reactant.promote_to(TracedRNumber{Int32}, 0),
-        ),
-        (),
-        "comparator",
-        false;
-        args_in_result=:none,
-        return_dialect=:stablehlo,
-    ).f
+    func =
+        Reactant.TracedUtils.make_mlir_fn(
+            comparator,
+            (
+                Reactant.promote_to(TracedRNumber{T}, 0),
+                Reactant.promote_to(TracedRNumber{T}, 0),
+                Reactant.promote_to(TracedRNumber{Int32}, 0),
+                Reactant.promote_to(TracedRNumber{Int32}, 0),
+            ),
+            (),
+            "comparator",
+            false;
+            args_in_result=:none,
+            return_dialect=:stablehlo,
+        ).f
     @assert MLIR.IR.nregions(func) == 1
     fn_name = MLIR.IR.FlatSymbolRefAttribute(
         String(
@@ -2178,15 +2182,16 @@ end
         Reactant.promote_to(TracedRNumber, zero(T)),
     )
 
-    compiled_fn = Reactant.TracedUtils.make_mlir_fn(
-        f,
-        sample_inputs,
-        (),
-        "update_computation",
-        false;
-        args_in_result=:result,
-        return_dialect=:stablehlo,
-    ).f
+    compiled_fn =
+        Reactant.TracedUtils.make_mlir_fn(
+            f,
+            sample_inputs,
+            (),
+            "update_computation",
+            false;
+            args_in_result=:result,
+            return_dialect=:stablehlo,
+        ).f
     update_computation = MLIR.IR.Region()
     MLIR.API.mlirRegionTakeBody(update_computation, MLIR.IR.region(compiled_fn, 1))
     MLIR.IR.rmfromparent!(compiled_fn)
@@ -2370,34 +2375,36 @@ end
 
     input_types = [mlir_type(arg) for arg in linear_args]
 
-    cond_fn_compiled = Reactant.TracedUtils.make_mlir_fn(
-        cond_fn,
-        traced_args,
-        (),
-        string(gensym("cond_fn")),
-        false;
-        return_dialect=:stablehlo,
-        args_in_result=:result,
-        do_transpose=false,
-        argprefix=gensym("loop_condarg"),
-        resprefix=gensym("loop_condres"),
-        resargprefix=gensym("loop_condresarg"),
-    ).f
-
-    body_fn_compiled = Reactant.TracedUtils.make_mlir_fn(
-        body_fn,
-        traced_args,
-        (),
-        string(gensym("body_fn")),
-        false;
-        return_dialect=:stablehlo,
-        args_in_result=:all,
-        do_transpose=false,
-        verify_arg_names,
-        argprefix=gensym("loop_bodyarg"),
-        resprefix=gensym("loop_bodyres"),
-        resargprefix=gensym("loop_bodyresarg"),
-    ).f
+    cond_fn_compiled =
+        Reactant.TracedUtils.make_mlir_fn(
+            cond_fn,
+            traced_args,
+            (),
+            string(gensym("cond_fn")),
+            false;
+            return_dialect=:stablehlo,
+            args_in_result=:result,
+            do_transpose=false,
+            argprefix=gensym("loop_condarg"),
+            resprefix=gensym("loop_condres"),
+            resargprefix=gensym("loop_condresarg"),
+        ).f
+
+    body_fn_compiled =
+        Reactant.TracedUtils.make_mlir_fn(
+            body_fn,
+            traced_args,
+            (),
+            string(gensym("body_fn")),
+            false;
+            return_dialect=:stablehlo,
+            args_in_result=:all,
+            do_transpose=false,
+            verify_arg_names,
+            argprefix=gensym("loop_bodyarg"),
+            resprefix=gensym("loop_bodyres"),
+            resargprefix=gensym("loop_bodyresarg"),
+        ).f
 
     cond_reg = Reactant.TracedUtils.__take_region(cond_fn_compiled)
     body_reg = Reactant.TracedUtils.__take_region(body_fn_compiled)
@@ -2842,17 +2849,16 @@ end
         location,
     )
 
-    corrected_traced_results = map(
-        zip(traced_false_results, traced_true_results)
-    ) do (fr, tr)
-        if fr isa MissingTracedValue && tr isa MissingTracedValue
-            return fr
-        elseif fr isa MissingTracedValue
-            return tr
-        else
-            return fr
+    corrected_traced_results =
+        map(zip(traced_false_results, traced_true_results)) do (fr, tr)
+            if fr isa MissingTracedValue && tr isa MissingTracedValue
+                return fr
+            elseif fr isa MissingTracedValue
+                return tr
+            else
+                return fr
+            end
         end
-    end
 
     @assert length(all_paths) == length(result_types)
 
@@ -3491,15 +3497,16 @@ end
 function _construct_reduce_function(f::F, Ts::Type...) where {F}
     inputs_1 = [Reactant.promote_to(TracedRNumber{T}, 0) for T in Ts]
     inputs_2 = [Reactant.promote_to(TracedRNumber{T}, 0) for T in Ts]
-    func = Reactant.TracedUtils.make_mlir_fn(
-        f,
-        (inputs_1..., inputs_2...),
-        (),
-        "reduce_fn" * string(f),
-        false;
-        args_in_result=:none,
-        return_dialect=:stablehlo,
-    ).f
+    func =
+        Reactant.TracedUtils.make_mlir_fn(
+            f,
+            (inputs_1..., inputs_2...),
+            (),
+            "reduce_fn" * string(f),
+            false;
+            args_in_result=:none,
+            return_dialect=:stablehlo,
+        ).f
 
     @assert MLIR.IR.nregions(func) == 1
     ftype_attr = MLIR.IR.getattr(func, "function_type")
diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
index 7517984f57..58eeed4f1f 100644
--- a/src/TracedRArray.jl
+++ b/src/TracedRArray.jl
@@ -1022,7 +1022,7 @@ end
         end
         return TracedRNumber{
             unwrapped_eltype(
-                Base._accumulate_promote_op(op, Array{T,ndims(A)}(undef, size(A));init)
+                Base._accumulate_promote_op(op, Array{T,ndims(A)}(undef, size(A)); init)
             ),
         }
     end
diff --git a/src/TracedUtils.jl b/src/TracedUtils.jl
index 596af5f5a6..b2d252f2c8 100644
--- a/src/TracedUtils.jl
+++ b/src/TracedUtils.jl
@@ -96,17 +96,17 @@ function ReactantCore.materialize_traced_array(
 end
 
 get_mlir_data(x::TracedRNumber) = x.mlir_data
-set_mlir_data!(x::TracedRNumber, data) = (x.mlir_data=data; return x)
+set_mlir_data!(x::TracedRNumber, data) = (x.mlir_data = data; return x)
 get_paths(x::TracedRNumber) = x.paths
-set_paths!(x::TracedRNumber, paths) = (x.paths=paths; return x)
+set_paths!(x::TracedRNumber, paths) = (x.paths = paths; return x)
 
 get_mlir_data(x::TracedRArray) = x.mlir_data
 get_mlir_data(x::AnyTracedRArray) = get_mlir_data(materialize_traced_array(x))
 get_paths(x::TracedRArray) = x.paths
-set_paths!(x::TracedRArray, paths) = (x.paths=paths; return x)
+set_paths!(x::TracedRArray, paths) = (x.paths = paths; return x)
 
 get_paths(x::MissingTracedValue) = x.paths
-set_paths!(x::MissingTracedValue, paths) = (x.paths=paths; return x)
+set_paths!(x::MissingTracedValue, paths) = (x.paths = paths; return x)
 
 function set_mlir_data!(x::TracedRArray, data)
     x.mlir_data = data
@@ -1294,11 +1294,12 @@ function elem_apply(f, args::Vararg{Any,Nargs}) where {Nargs}
         push_val!(batch_inputs, ogarg, path[3:end])
 
         if ogarg isa Base.RefValue
-            batch_inputs[end] = (@opcall broadcast_in_dim(
-                TracedRArray(batch_inputs[end]),
-                Int64[],
-                collect(Int64, input_shapes[1]),
-            )).mlir_data
+            batch_inputs[end] =
+                (@opcall broadcast_in_dim(
+                    TracedRArray(batch_inputs[end]),
+                    Int64[],
+                    collect(Int64, input_shapes[1]),
+                )).mlir_data
         end
     end
 
diff --git a/src/Types.jl b/src/Types.jl
index bf2b88e1a4..2fd2a7df9e 100644
--- a/src/Types.jl
+++ b/src/Types.jl
@@ -781,11 +781,13 @@ function InterpolateArray(
                         dim -> c[dim] == 1 ? lows[dim][I.I[dim]] : highs[dim][I.I[dim]], N
                     )
 
-                    w_int = prod(ntuple(dim -> if c[dim] == 1
-                        (dens[dim] - rems[dim][I.I[dim]])
-                    else
-                        rems[dim][I.I[dim]]
-                    end, N))
+                    w_int = prod(
+                        ntuple(dim -> if c[dim] == 1
+                            (dens[dim] - rems[dim][I.I[dim]])
+                        else
+                            rems[dim][I.I[dim]]
+                        end, N)
+                    )
 
                     sum_val += w_int * local_cpu_array[CartesianIndex(idx)]
                 end
diff --git a/src/compiler/Codegen.jl b/src/compiler/Codegen.jl
index 996df09d32..788fdc1143 100644
--- a/src/compiler/Codegen.jl
+++ b/src/compiler/Codegen.jl
@@ -247,9 +247,12 @@ Base.@nospecializeinfer function create_result(
 
         result = Expr(:new, T, elems...)
 
-        push!(resultgen_code, quote
-            $sym = $result
-        end)
+        push!(
+            resultgen_code,
+            quote
+                $sym = $result
+            end,
+        )
         result_cache[tocopy] = sym
     end
 
@@ -275,9 +278,12 @@ Base.@nospecializeinfer function create_result(
 
         result = Meta.quot(tocopy)
 
-        push!(resultgen_code, quote
-            $sym = $result
-        end)
+        push!(
+            resultgen_code,
+            quote
+                $sym = $result
+            end,
+        )
         result_cache[tocopy] = sym
     end
 
@@ -314,9 +320,12 @@ function create_pjrt_number_result(
         else
             result = :(ConcretePJRTNumber{$T}($restore))
         end
-        push!(resultgen_code, quote
-            $sym = $result
-        end)
+        push!(
+            resultgen_code,
+            quote
+                $sym = $result
+            end,
+        )
         result_cache[tocopy] = sym
     end
 
@@ -353,9 +362,12 @@ function create_ifrt_number_result(
         else
             result = :(ConcreteIFRTNumber{$T}($restore))
         end
-        push!(resultgen_code, quote
-            $sym = $result
-        end)
+        push!(
+            resultgen_code,
+            quote
+                $sym = $result
+            end,
+        )
         result_cache[tocopy] = sym
     end
 
@@ -435,9 +447,12 @@ function create_result(
         else
             result = :(ConcretePJRTArray{$T,$N}($restore, $(tocopy.shape)))
         end
-        push!(resultgen_code, quote
-            $sym = $result
-        end)
+        push!(
+            resultgen_code,
+            quote
+                $sym = $result
+            end,
+        )
         result_cache[tocopy] = sym
     end
 
@@ -491,9 +506,12 @@ function create_result(
         else
             result = :(ConcreteIFRTArray{$T,$N}($(restore), $(tocopy.shape)))
         end
-        push!(resultgen_code, quote
-            $sym = $result
-        end)
+        push!(
+            resultgen_code,
+            quote
+                $sym = $result
+            end,
+        )
         result_cache[tocopy] = sym
     end
 
@@ -556,17 +574,23 @@ function create_result(
         sym = Symbol("result", var_idx[])
         var_idx[] += 1
 
-        push!(resultgen_code, quote
-            $sym = $(Array{T,N})(undef, $(size(tocopy)...,))
-        end)
+        push!(
+            resultgen_code,
+            quote
+                $sym = $(Array{T,N})(undef, $(size(tocopy)...,))
+            end,
+        )
 
         result_cache[tocopy] = sym
 
         for (i, v) in enumerate(tocopy)
             subexpr = create_result(v, append_path(path, i), args...)
-            push!(resultgen_code, quote
-                @inbounds $sym[$i] = $subexpr
-            end)
+            push!(
+                resultgen_code,
+                quote
+                    @inbounds $sym[$i] = $subexpr
+                end,
+            )
         end
     end
 
@@ -664,9 +688,12 @@ function create_result(
         sym = Symbol("result", var_idx[])
         var_idx[] += 1
 
-        push!(resultgen_code, quote
-            $sym = $D()
-        end)
+        push!(
+            resultgen_code,
+            quote
+                $sym = $D()
+            end,
+        )
 
         result_cache[tocopy] = sym
 
@@ -675,9 +702,12 @@ function create_result(
             # symbol keys must be quoted in generated code; otherwise
             # they are interpreted as variable references
             k_expr = k isa Symbol ? QuoteNode(k) : k
-            push!(resultgen_code, quote
-                @inbounds $sym[$k_expr] = $subexpr
-            end)
+            push!(
+                resultgen_code,
+                quote
+                    @inbounds $sym[$k_expr] = $subexpr
+                end,
+            )
         end
     end
 
@@ -1231,9 +1261,12 @@ function codegen_unflatten!(
             sym = Symbol("result", var_idx[])
             var_idx[] += 1
 
-            push!(resultgen_code, quote
-                $sym = $argres.data
-            end)
+            push!(
+                resultgen_code,
+                quote
+                    $sym = $argres.data
+                end,
+            )
 
             result_stores[path] = sym
         end
diff --git a/test/core/compile.jl b/test/core/compile.jl
index da6fd20cf2..5eddd4eeab 100644
--- a/test/core/compile.jl
+++ b/test/core/compile.jl
@@ -236,9 +236,12 @@ end
             sym = Symbol("result", var_idx[])
             var_idx[] += 1
 
-            push!(resultgen_code, quote
-                $sym = ($MockTestCustomPath)($ar)
-            end)
+            push!(
+                resultgen_code,
+                quote
+                    $sym = ($MockTestCustomPath)($ar)
+                end,
+            )
             result_cache[tocopy] = sym
         end
 
diff --git a/test/integration/structarrays.jl b/test/integration/structarrays.jl
index 69eb884be3..3ce8a41e21 100644
--- a/test/integration/structarrays.jl
+++ b/test/integration/structarrays.jl
@@ -9,7 +9,7 @@ using StructArrays, StaticArrays, Reactant, LinearAlgebra, Test
     # Note that the element type (the NamedTuple) contains ConcreteRNumbers even though track_numbers is not enabled.
     # This is because when the backing arrays are converted to TracedRArrays, their elements will contain TracedRNumbers.
     # In order for the element type to match the backing arrays, we need to use ConcreteRNumbers here as well:
-    @test typeof(x_ra)<:(StructArray{
+    @test typeof(x_ra) <: (StructArray{
         <:NamedTuple{
             (:a, :b, :c),
             <:Tuple{<:ConcreteRNumber{Float64},String,<:ConcreteRNumber{Float32}},
@@ -26,7 +26,7 @@ using StructArrays, StaticArrays, Reactant, LinearAlgebra, Test
         Reactant.make_tracer(Reactant.OrderedIdDict(), x_ra, (), Reactant.ConcreteToTraced)
     ) == StructArray{
         @NamedTuple{
-            a::Reactant.TracedRFloat{Float64},b::String,c::Reactant.TracedRFloat{Float32}
+            a::Reactant.TracedRFloat{Float64}, b::String, c::Reactant.TracedRFloat{Float32}
         },
         2,
         @NamedTuple{

From 0210c814f6bd82e8eb895630e758b4c018c00262 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Wed, 8 Jul 2026 21:13:52 +0200
Subject: [PATCH 19/30] Drop formatter churn unrelated to the number-type split

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.
---
 ext/ReactantAbstractFFTsExt.jl       |  8 ++++----
 ext/ReactantCUDAExt.jl               |  2 +-
 ext/ReactantKernelAbstractionsExt.jl |  3 ++-
 ext/ReactantMPIExt/Ops.jl            |  4 ++--
 src/stdlibs/LinearAlgebra.jl         | 14 ++++++--------
 5 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/ext/ReactantAbstractFFTsExt.jl b/ext/ReactantAbstractFFTsExt.jl
index 417ee694cc..335dfeeda6 100644
--- a/ext/ReactantAbstractFFTsExt.jl
+++ b/ext/ReactantAbstractFFTsExt.jl
@@ -135,7 +135,7 @@ for op in (:rfft, :fft, :ifft)
         if __is_valid_stablehlo_fft_dims(dims, ndims(x))
             return @opcall fft(
                 TracedUtils.materialize_traced_array(x);
-                type=($(uppercase(string(op)))),
+                type=$(uppercase(string(op))),
                 length=fft_lengths,
             )
         end
@@ -144,7 +144,7 @@ for op in (:rfft, :fft, :ifft)
             @opcall(
                 fft(
                     TracedUtils.materialize_traced_array(permutedims(x, perm));
-                    type=($(uppercase(string(op)))),
+                    type=$(uppercase(string(op))),
                     length=fft_lengths,
                 )
             ),
@@ -233,7 +233,7 @@ for op in (:irfft,)
         if __is_valid_stablehlo_fft_dims(dims, ndims(x))
             return @opcall fft(
                 TracedUtils.materialize_traced_array(x);
-                type=($(uppercase(string(op)))),
+                type=$(uppercase(string(op))),
                 length=fft_lengths,
             )
         end
@@ -243,7 +243,7 @@ for op in (:irfft,)
             @opcall(
                 fft(
                     TracedUtils.materialize_traced_array(permutedims(x, perm));
-                    type=($(uppercase(string(op)))),
+                    type=$(uppercase(string(op))),
                     length=fft_lengths,
                 )
             ),
diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl
index 5aae5f9d37..bc27680e74 100644
--- a/ext/ReactantCUDAExt.jl
+++ b/ext/ReactantCUDAExt.jl
@@ -1572,7 +1572,7 @@ const ReactantCUDAJob = GPUCompiler.CompilerJob{
 GPUCompiler.can_vectorize(job::ReactantCUDAJob) = !job.config.params.raising
 function GPUCompiler.optimization_options(job::ReactantCUDAJob)
     raising = job.config.params.raising
-    return (; instcombine=(!raising), fastmath=(!raising), aggressiveinstcombine=(!raising))
+    return (; instcombine=!raising, fastmath=!raising, aggressiveinstcombine=!raising)
 end
 
 function GPUCompiler.method_table(@nospecialize(job::ReactantCUDAJob))
diff --git a/ext/ReactantKernelAbstractionsExt.jl b/ext/ReactantKernelAbstractionsExt.jl
index fa9ccaccf0..954b153c92 100644
--- a/ext/ReactantKernelAbstractionsExt.jl
+++ b/ext/ReactantKernelAbstractionsExt.jl
@@ -112,8 +112,9 @@ end
 @static if VERSION < v"1.12-"
     Reactant.@reactant_overlay Base.@nospecializeinfer @noinline function (
         obj::KA.Kernel{ReactantBackend}
+    )(
+        @nospecialize args...; ndrange=nothing, workgroupsize=nothing
     )
-        (@nospecialize args...; ndrange=nothing, workgroupsize=nothing)
         return Reactant.call_with_reactant(
             Reactant.ka_with_reactant, ndrange, workgroupsize, obj, args...
         )
diff --git a/ext/ReactantMPIExt/Ops.jl b/ext/ReactantMPIExt/Ops.jl
index 4bf0e04ce4..88eb0db359 100644
--- a/ext/ReactantMPIExt/Ops.jl
+++ b/ext/ReactantMPIExt/Ops.jl
@@ -225,7 +225,7 @@ const MPI_OP_MAP = Dict(
 
 function get_mpi_op_enum(op)
     return get(MPI_OP_MAP, op.val) do
-        return throw(ArgumentError("Unknown MPI op `$op`"))
+        throw(ArgumentError("Unknown MPI op `$op`"))
     end
 end
 
@@ -261,7 +261,7 @@ const MPI_DATATYPE_MAP = Dict(
 
 function get_mpi_datatype_enum(datatype)
     return get(MPI_DATATYPE_MAP, datatype.val) do
-        return throw(ArgumentError("Unknown MPI datatype `$datatype`"))
+        throw(ArgumentError("Unknown MPI datatype `$datatype`"))
     end
 end
 
diff --git a/src/stdlibs/LinearAlgebra.jl b/src/stdlibs/LinearAlgebra.jl
index 2c7bbb2ca9..1c4b343763 100644
--- a/src/stdlibs/LinearAlgebra.jl
+++ b/src/stdlibs/LinearAlgebra.jl
@@ -118,7 +118,7 @@ for (AT, comp) in ((:LowerTriangular, "GE"), (:UpperTriangular, "LE"))
             px = materialize_traced_array(parent(x))
             row_idxs = @opcall iota(Int, [m, n]; iota_dimension=1)
             col_idxs = @opcall iota(Int, [m, n]; iota_dimension=2)
-            indicator = @opcall compare(row_idxs, col_idxs; comparison_direction=($(comp)))
+            indicator = @opcall compare(row_idxs, col_idxs; comparison_direction=$(comp))
             return @opcall select(indicator, px, zero(px))
         end
 
@@ -214,11 +214,9 @@ for (AT, dcomp, ocomp) in (
         m, n = size(x)
         row_idxs = @opcall iota(Int, [m, n]; iota_dimension=1)
         col_idxs = @opcall iota(Int, [m, n]; iota_dimension=2)
-        data_indicator = @opcall compare(
-            row_idxs, col_idxs; comparison_direction=($(dcomp))
-        )
+        data_indicator = @opcall compare(row_idxs, col_idxs; comparison_direction=$(dcomp))
         original_indicator = @opcall compare(
-            row_idxs, col_idxs; comparison_direction=($(ocomp))
+            row_idxs, col_idxs; comparison_direction=$(ocomp)
         )
         res = @opcall add(
             @opcall(select(data_indicator, tdata, z)),
@@ -786,9 +784,9 @@ for (wT, lower, ud) in (
             parent(A),
             rhs;
             left_side=false,
-            lower=($(lower)),
-            transpose_a=('N'),
-            unit_diagonal=($(ud)),
+            lower=$(lower),
+            transpose_a='N',
+            unit_diagonal=$(ud),
         )
     end
 end

From 362cf8fed7862254a2ddb017e1d53db61b6b92f0 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 9 Jul 2026 10:09:07 +0200
Subject: [PATCH 20/30] Exclude `Union{}` element types from `AnyTracedRArray`

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 (#1290). This restores the
upstream subtype guard test alongside the behavior tests.

Created by generative AI.
---
 src/Types.jl         | 12 +++++++++++-
 test/core/tracing.jl |  8 ++++----
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/Types.jl b/src/Types.jl
index 2fd2a7df9e..a654d3d458 100644
--- a/src/Types.jl
+++ b/src/Types.jl
@@ -233,7 +233,17 @@ end
 @leaf TracedRational
 Adapt.parent_type(::Type{TracedRational{T}}) where {T} = TracedRational{T}
 
-const AnyTracedRArray{T,N} = AbstractArray{TracedRNumber{T},N}
+## Enumerate the traced element types invariantly instead of bounding them
+## covariantly (`AbstractArray{<:TracedRNumber{T},N}`): a covariant bound would
+## also match `Vector{Union{}}` (the type of empty array literals), which must
+## not be routed into the traced-array machinery (#1290).
+const AnyTracedRArray{T,N} = Union{
+    AbstractArray{TracedRInteger{T},N},
+    AbstractArray{TracedRFloat{T},N},
+    AbstractArray{TracedRComplex{T},N},
+    AbstractArray{TracedRReal{T},N},
+    AbstractArray{TracedRNumber{T},N},
+}
 const AnyTracedRVector{T} = AnyTracedRArray{T,1}
 const AnyTracedRMatrix{T} = AnyTracedRArray{T,2}
 const AnyTracedRVecOrMat{T} = Union{AnyTracedRVector{T},AnyTracedRMatrix{T}}
diff --git a/test/core/tracing.jl b/test/core/tracing.jl
index 32a4b63ec7..fe0717eb82 100644
--- a/test/core/tracing.jl
+++ b/test/core/tracing.jl
@@ -26,10 +26,10 @@ struct RMSProp{Teta,Trho,Teps,C<:Bool}
 end
 
 @testset "Traced Type" begin
-    # `Vector{Union{}}` (e.g. an empty array literal) matches the covariant
-    # `AnyTracedRArray` alias, but must not be routed into the traced-array
-    # machinery
-    @test ReactantCore.materialize_traced_array(Union{}[]) isa Vector{Union{}}
+    # `Vector{Union{}}` (e.g. an empty array literal) must not match the
+    # `AnyTracedRArray` alias or be routed into the traced-array machinery
+    @test !(Vector{Union{}} <: Reactant.AnyTracedRArray)
+    @test Reactant.materialize_traced_array(Union{}[]) isa Vector{Union{}}
     @test Reactant.aos_to_soa(Union{}[]) isa Vector{Union{}}
 end
 

From e6b56d7688ea4059ed83c27d1e0627d66b88ed2e Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 9 Jul 2026 10:09:21 +0200
Subject: [PATCH 21/30] Expect kind-typed traced arrays in the tracing tests

`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.
---
 test/core/tracing.jl | 79 ++++++++++++++++++++++++++++----------------
 1 file changed, 50 insertions(+), 29 deletions(-)

diff --git a/test/core/tracing.jl b/test/core/tracing.jl
index fe0717eb82..be4268bdc8 100644
--- a/test/core/tracing.jl
+++ b/test/core/tracing.jl
@@ -2,6 +2,7 @@ using Reactant, Test
 using Reactant:
     traced_type,
     TracedRArray,
+    TracedRFloat,
     TracedRNumber,
     ConcreteToTraced,
     ArrayToConcrete,
@@ -80,17 +81,33 @@ end
         (Complex{UInt64}, Complex{UInt64}, TracedRNumber{Complex{UInt64}}),
 
         # RArray types
-        (ConcreteRArray{Float64,0}, TracedRArray{Float64,0}, TracedRArray{Float64,0}),
-        (ConcreteRArray{Float64,1}, TracedRArray{Float64,1}, TracedRArray{Float64,1}),
-        (ConcreteRArray{Float64,2}, TracedRArray{Float64,2}, TracedRArray{Float64,2}),
-        (ConcreteRArray{Float64,3}, TracedRArray{Float64,3}, TracedRArray{Float64,3}),
+        (
+            ConcreteRArray{Float64,0},
+            TracedRArray{Float64,0,TracedRFloat{Float64}},
+            TracedRArray{Float64,0,TracedRFloat{Float64}},
+        ),
+        (
+            ConcreteRArray{Float64,1},
+            TracedRArray{Float64,1,TracedRFloat{Float64}},
+            TracedRArray{Float64,1,TracedRFloat{Float64}},
+        ),
+        (
+            ConcreteRArray{Float64,2},
+            TracedRArray{Float64,2,TracedRFloat{Float64}},
+            TracedRArray{Float64,2,TracedRFloat{Float64}},
+        ),
+        (
+            ConcreteRArray{Float64,3},
+            TracedRArray{Float64,3,TracedRFloat{Float64}},
+            TracedRArray{Float64,3,TracedRFloat{Float64}},
+        ),
 
         # Array types
         (Array{Float64,1}, Array{Float64,1}, Array{TracedRNumber{Float64},1}),
         (
             Array{ConcreteRArray{Float64,2},1},
-            Array{TracedRArray{Float64,2},1},
-            Array{TracedRArray{Float64,2},1},
+            Array{TracedRArray{Float64,2,TracedRFloat{Float64}},1},
+            Array{TracedRArray{Float64,2,TracedRFloat{Float64}},1},
         ),
 
         # AbstractArray types
@@ -122,8 +139,8 @@ end
         (Union{Nothing,Int}, Union{Nothing,Int}, Union{Nothing,TracedRNumber{Int}}),
         (
             Union{Nothing,ConcreteRArray{Float64,1}},
-            Union{Nothing,TracedRArray{Float64,1}},
-            Union{Nothing,TracedRArray{Float64,1}},
+            Union{Nothing,TracedRArray{Float64,1,TracedRFloat{Float64}}},
+            Union{Nothing,TracedRArray{Float64,1,TracedRFloat{Float64}}},
         ),
 
         # UnionAll types
@@ -134,7 +151,11 @@ end
         (Array{>:Float32}, Array{>:Float32}, Array{>:Float32}),
         (ConcreteRNumber, TracedRNumber, TracedRNumber),
         (ConcreteRArray, TracedRArray, TracedRArray),
-        (ConcreteRArray{Float64}, TracedRArray{Float64}, TracedRArray{Float64}),
+        (
+            ConcreteRArray{Float64},
+            (TracedRArray{Float64,N,TracedRFloat{Float64}} where {N}),
+            (TracedRArray{Float64,N,TracedRFloat{Float64}} where {N}),
+        ),
         (
             ConcreteRArray{T,2} where {T},
             TracedRArray{T,2} where {T},
@@ -159,8 +180,8 @@ end
         ),
         (
             Union{Nothing,ConcreteRArray{Float64}},
-            Union{Nothing,TracedRArray{Float64}},
-            Union{Nothing,TracedRArray{Float64}},
+            Union{Nothing,(TracedRArray{Float64,N,TracedRFloat{Float64}} where {N})},
+            Union{Nothing,(TracedRArray{Float64,N,TracedRFloat{Float64}} where {N})},
         ),
         (
             Union{Nothing,ConcreteRArray{T,2} where {T}},
@@ -177,8 +198,8 @@ end
         (Ptr{Float64}, Ptr{Float64}, Ptr{TracedRNumber{Float64}}),
         (
             Ptr{ConcreteRArray{Float64,1}},
-            Ptr{TracedRArray{Float64,1}},
-            Ptr{TracedRArray{Float64,1}},
+            Ptr{TracedRArray{Float64,1,TracedRFloat{Float64}}},
+            Ptr{TracedRArray{Float64,1,TracedRFloat{Float64}}},
         ),
         (
             Core.LLVMPtr{Float64},
@@ -187,8 +208,8 @@ end
         ),
         (
             Core.LLVMPtr{ConcreteRArray{Float64,1}},
-            Core.LLVMPtr{TracedRArray{Float64,1}},
-            Core.LLVMPtr{TracedRArray{Float64,1}},
+            Core.LLVMPtr{TracedRArray{Float64,1,TracedRFloat{Float64}}},
+            Core.LLVMPtr{TracedRArray{Float64,1,TracedRFloat{Float64}}},
         ),
         (
             Base.RefValue{Float64},
@@ -197,21 +218,21 @@ end
         ),
         (
             Base.RefValue{ConcreteRArray{Float64,1}},
-            Base.RefValue{TracedRArray{Float64,1}},
-            Base.RefValue{TracedRArray{Float64,1}},
+            Base.RefValue{TracedRArray{Float64,1,TracedRFloat{Float64}}},
+            Base.RefValue{TracedRArray{Float64,1,TracedRFloat{Float64}}},
         ),
         (Ref{Float64}, Ref{Float64}, Ref{TracedRNumber{Float64}}),
         (
             Ref{ConcreteRArray{Float64,1}},
-            Ref{TracedRArray{Float64,1}},
-            Ref{TracedRArray{Float64,1}},
+            Ref{TracedRArray{Float64,1,TracedRFloat{Float64}}},
+            Ref{TracedRArray{Float64,1,TracedRFloat{Float64}}},
         ),
 
         # Function types
         (
             MyFix{2,typeof(mul),ConcreteRArray{Float64,1}},
-            MyFix{2,typeof(mul),TracedRArray{Float64,1}},
-            MyFix{2,typeof(mul),TracedRArray{Float64,1}},
+            MyFix{2,typeof(mul),TracedRArray{Float64,1,TracedRFloat{Float64}}},
+            MyFix{2,typeof(mul),TracedRArray{Float64,1,TracedRFloat{Float64}}},
         ),
 
         # Val types
@@ -220,8 +241,8 @@ end
         (Val{:x}, Val{:x}, Val{:x}),
         (
             Dict{Int,ConcreteRArray{Float64,0}},
-            Dict{Int,TracedRArray{Float64,0}},
-            Dict{Int,TracedRArray{Float64,0}},
+            Dict{Int,TracedRArray{Float64,0,TracedRFloat{Float64}}},
+            Dict{Int,TracedRArray{Float64,0,TracedRFloat{Float64}}},
         ),
 
         # others
@@ -229,8 +250,8 @@ end
         (Dict, Dict, Dict),
         (
             (Dict{A,ConcreteRArray{Float64,0}} where {A}),
-            (Dict{A,TracedRArray{Float64,0}} where {A}),
-            (Dict{A,TracedRArray{Float64,0}} where {A}),
+            (Dict{A,TracedRArray{Float64,0,TracedRFloat{Float64}}} where {A}),
+            (Dict{A,TracedRArray{Float64,0,TracedRFloat{Float64}}} where {A}),
         ),
         (
             (
@@ -273,15 +294,15 @@ end
         ),
         (
             Wrapper{Float64,ConcreteRArray{Float64,1}},
-            Wrapper{Float64,TracedRArray{Float64,1}},
-            Wrapper{TracedRNumber{Float64},TracedRArray{Float64,1}},
+            Wrapper{Float64,TracedRArray{Float64,1,TracedRFloat{Float64}}},
+            Wrapper{TracedRNumber{Float64},TracedRArray{Float64,1,TracedRFloat{Float64}}},
         ),
         (Wrapper{Symbol}, Wrapper{Symbol}, Wrapper{Symbol}),
         (Wrapper{Float64}, Wrapper{Float64}, Wrapper{TracedRNumber{Float64}}),
         (
             Wrapper{ConcreteRArray{Float64,1}},
-            Wrapper{TracedRArray{Float64,1}},
-            Wrapper{TracedRArray{Float64,1}},
+            Wrapper{TracedRArray{Float64,1,TracedRFloat{Float64}}},
+            Wrapper{TracedRArray{Float64,1,TracedRFloat{Float64}}},
         ),
         (Wrapper, Wrapper, Wrapper),
     ]

From 07a74a81cbc716017fd160c55494a0f828eb2e68 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 9 Jul 2026 10:09:36 +0200
Subject: [PATCH 22/30] Keep the `apply_type_with_promotion` docstring attached

`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.
---
 src/Tracing.jl | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/Tracing.jl b/src/Tracing.jl
index f9da34f986..e115c628d7 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -968,6 +968,14 @@ resolve_conflict(t1::Type{<:ConcreteRNumber{T}}, t2::Type{T}) where {T} = T
 resolve_conflict(t1::Type{T}, t2::Type{<:ConcreteRNumber{T}}) where {T} = T
 resolve_conflict(t1, t2) = promote_type(t1, t2)
 
+# `TracedRArray{T,N}` with the element-type parameter left free has no single
+# eltype and hence cannot instantiate an `AbstractArray{T}`-bounded type
+# parameter; substitute the concrete traced array type.
+concretize_traced_param(@nospecialize(T)) = T
+function concretize_traced_param(::Type{TracedRArray{T,N}}) where {T,N}
+    return TracedRArray{T,N,traced_number_type(T)}
+end
+
 """
 This function tries to apply the param types to the wrapper type.
 When there's a constraint conflict, it tries to resolve it:
@@ -982,23 +990,15 @@ using Reactant
 struct Foo{T, A<:AbstractArray{T}}
     a::A
 end
-Reactant.apply_type_with_promotion(Foo, (Int, TracedRArray{Int, 1}))
+Reactant.apply_type_with_promotion(Foo, Any[Int, TracedRArray{Int, 1}])
 ```
 returns
 ```jl
-(Foo{Reactant.TracedRInteger{Int64}, Reactant.TracedRArray{Int64, 1}}, Bool[1, 0])
+(Foo{Reactant.TracedRInteger{Int64}, Reactant.TracedRArray{Int64, 1, Reactant.TracedRInteger{Int64}}}, Bool[1, 1])
 ```
 
 The first type parameter has been promoted to satisfy to be in agreement with the second parameter.
 """
-# `TracedRArray{T,N}` with the element-type parameter left free has no single
-# eltype and hence cannot instantiate an `AbstractArray{T}`-bounded type
-# parameter; substitute the concrete traced array type.
-concretize_traced_param(@nospecialize(T)) = T
-function concretize_traced_param(::Type{TracedRArray{T,N}}) where {T,N}
-    return TracedRArray{T,N,traced_number_type(T)}
-end
-
 function apply_type_with_promotion(wrapper, params, relevant_typevars=typevar_dict(wrapper))
     unwrapped = Base.unwrap_unionall(wrapper) # remove all the typevars
     original_params = copy(params)

From 57c1a9013562319fe67c5d81b4aec91b4a854561 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 9 Jul 2026 10:09:36 +0200
Subject: [PATCH 23/30] Disambiguate `UnitRange` `getindex` against 1.10's
 overflow-safe grouping

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.
---
 src/Indexing.jl | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/Indexing.jl b/src/Indexing.jl
index 2b71a60ac2..78debe9584 100644
--- a/src/Indexing.jl
+++ b/src/Indexing.jl
@@ -96,6 +96,15 @@ for (aT, iT) in (
     (Base.IdentityUnitRange, TracedRInteger),
     # 1.10 specific ambiguity fixes
     (UnitRange{<:TracedRNumber}, Int),
+    # 1.10's `getindex(::UnitRange{<:Union{...}}, ::Integer)` overflow-safe
+    # grouping (range.jl) needs an exact-grouping disambiguator now that
+    # traced integers are `Integer`s
+    (
+        UnitRange{
+            <:Union{Bool,Int128,Int16,Int32,Int64,Int8,UInt128,UInt16,UInt32,UInt64,UInt8}
+        },
+        TracedRInteger,
+    ),
 )
     @eval function Base.getindex(a::$aT, index::$iT)
         return convert(

From 9b59c5107d99ede26128dc48c0a3b68cc63d3edf Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 9 Jul 2026 10:13:44 +0200
Subject: [PATCH 24/30] Require Enzyme 0.13.177

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.
---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index b0020d4420..e0b1f364e6 100644
--- a/Project.toml
+++ b/Project.toml
@@ -118,7 +118,7 @@ DLFP8Types = "0.1"
 Dates = "1.10"
 Downloads = "1.6"
 EnumX = "1"
-Enzyme = "0.13.115"
+Enzyme = "0.13.177"
 EnzymeCore = "0.8.14"
 FFTW = "1"
 FileWatching = "1.10"

From 39f481dad1132d4bb841b7c2d81fa10aea501873 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 9 Jul 2026 16:27:35 +0200
Subject: [PATCH 25/30] Resolve getindex ambiguities against Base,
 LinearAlgebra, and SparseArrays

Now that traced integers are `Integer`s, the traced `getindex` methods
compete with the stdlib integer-indexing methods, most visibly on Julia
1.10 where those take `Integer` (1.12 takes `Int`). Match the competing
method shapes exactly:

- Dispatch the wrapped-array scalar `getindex` per traced element type:
  with the whole `AnyTracedRArray` union in one signature, specificity
  against fixed-arity structured-matrix methods is undecidable.
- Route scalar range indexing through one `traced_range_getindex`
  implementation per range family, with Base-facing shims covering each
  specialized range type against each traced index shape; the shims bind
  the element type as a method typevar like Base's methods do, and group
  `LinRange`/`StepRangeLen` like Base's own union-shaped method.
- Give the structured-matrix covers in the LinearAlgebra overlay both
  index shapes: exact `Int` pairs (1.12's methods) and
  `Union{Int,TracedRInteger{Int}}` (1.10's `Integer` methods).
- Shape the `ReadOnly` covers in the SparseArrays extension per traced
  element type with invariant storage, replacing the covariant ones;
  mixed variance made concrete `ReadOnly` indexing ambiguous at runtime.

Also fixes `LinRange` `getindex` with a traced index, which cannot go
through `Base.lerpi` (it casts through the element type); interpolate
with traced arithmetic instead, and test traced range indexing.

Created by generative AI.
---
 ext/ReactantSparseArraysExt/ReadOnly.jl |  69 +++++++------
 src/Indexing.jl                         | 130 +++++++++++++-----------
 src/stdlibs/LinearAlgebra.jl            |  21 +++-
 test/core/ranges.jl                     |   7 ++
 4 files changed, 127 insertions(+), 100 deletions(-)

diff --git a/ext/ReactantSparseArraysExt/ReadOnly.jl b/ext/ReactantSparseArraysExt/ReadOnly.jl
index 18685535cd..8338fa8751 100644
--- a/ext/ReactantSparseArraysExt/ReadOnly.jl
+++ b/ext/ReactantSparseArraysExt/ReadOnly.jl
@@ -10,13 +10,6 @@ function Base.getindex(
     return getindex(parent(A), :)
 end
 
-function Base.getindex(
-    A::ReadOnly{<:Reactant.TracedRNumber{T},N,V},
-    i::Union{Integer,Reactant.TracedRNumber{<:Integer}},
-) where {T,N,V<:AbstractArray{<:Reactant.TracedRNumber{T},N}}
-    return getindex(parent(A), i)
-end
-
 function Base.getindex(
     A::ReadOnly{<:Reactant.TracedRNumber{T},N,V}, inds::AbstractArray
 ) where {T,N,V<:AbstractArray{<:Reactant.TracedRNumber{T},N}}
@@ -41,32 +34,38 @@ function Base.getindex(
     return getindex(parent(A), :)
 end
 
-function Base.getindex(
-    x::ReadOnly{<:TracedRNumber{T},N},
-    idx::Vararg{Union{Integer,Reactant.TracedRNumber{<:Integer}},N},
-) where {T,N}
-    return getindex(parent(x), idx...)
-end
-
-# Disambiguates against the generic traced-array `getindex`.
-function Base.getindex(
-    x::ReadOnly{<:TracedRNumber{T},N}, idx::Vararg{Union{Int,Reactant.TracedRNumber{Int}},N}
-) where {T,N}
-    return getindex(parent(x), idx...)
-end
-
-function Base.getindex(
-    x::SparseArrays.ReadOnly{<:Reactant.TracedRNumber{T},1,V},
-    idx::Union{Int64,Reactant.TracedRNumber{Int64}},
-) where {T,V<:AbstractArray{<:Reactant.TracedRNumber{T},1}}
-    return getindex(parent(x), idx)
-end
-
-function Base.getindex(
-    x::SparseArrays.ReadOnly{
-        Reactant.TracedRNumber{T},N,V
-    } where {V<:AbstractArray{<:Reactant.TracedRNumber{T},N}},
-    idx::Vararg{Union{Int64,Reactant.TracedRNumber{Int64}},N},
-) where {T,N}
-    return getindex(parent(x), idx...)
+# ReadOnly's own `getindex` just forwards to the parent, but it ties with the
+# traced-array `getindex` methods; these covers forward the same way and are
+# shaped per traced element type with invariant storage, exactly like the
+# traced-array methods they disambiguate against.
+for ET in (
+    Reactant.TracedRInteger,
+    Reactant.TracedRFloat,
+    Reactant.TracedRComplex,
+    Reactant.TracedRReal,
+    Reactant.TracedRNumber,
+)
+    @eval begin
+        function Base.getindex(
+            x::ReadOnly{$ET{T},N,V},
+            idx::Vararg{Union{Integer,Reactant.TracedRNumber{<:Integer}},N},
+        ) where {T,N,V<:AbstractArray{$ET{T},N}}
+            return getindex(parent(x), idx...)
+        end
+        function Base.getindex(
+            x::ReadOnly{$ET{T},N,V}, idx::Union{Integer,Reactant.TracedRNumber{<:Integer}}
+        ) where {T,N,V<:AbstractArray{$ET{T},N}}
+            return getindex(parent(x), idx)
+        end
+        function Base.getindex(
+            x::ReadOnly{$ET{T},N,V}, idx::Vararg{Union{Int,Reactant.TracedRNumber{Int}},N}
+        ) where {T,N,V<:AbstractArray{$ET{T},N}}
+            return getindex(parent(x), idx...)
+        end
+        function Base.getindex(
+            x::ReadOnly{$ET{T},1,V}, idx::Union{Int,Reactant.TracedRNumber{Int}}
+        ) where {T,V<:AbstractVector{$ET{T}}}
+            return getindex(parent(x), idx)
+        end
+    end
 end
diff --git a/src/Indexing.jl b/src/Indexing.jl
index 78debe9584..e314572b2a 100644
--- a/src/Indexing.jl
+++ b/src/Indexing.jl
@@ -4,7 +4,10 @@ using ..Reactant:
     Reactant,
     TracedRArray,
     TracedRNumber,
+    TracedRReal,
     TracedRInteger,
+    TracedRFloat,
+    TracedRComplex,
     TracedStepRangeLen,
     TracedUnitRange
 using ..Reactant: AnyTracedRArray, ancestor, unwrapped_eltype
@@ -74,11 +77,17 @@ function Base.getindex(a::TracedRArray{T,N}, indices::Vararg{Any,N}) where {T,N}
 end
 
 ### Wrapped Array Types
-function Base.getindex(
-    a::AnyTracedRArray{T,N}, index::Vararg{Union{Int,TracedRNumber{Int}},N}
-) where {T,N}
-    ancestor, idxs = TracedUtils.get_ancestor_and_indices(a, index...)
-    return getindex(ancestor, idxs...)
+## Defined per traced element type rather than on `AnyTracedRArray`: with the
+## whole `Union` in the signature, specificity against the fixed-arity
+## structured-matrix `getindex` methods in LinearAlgebra is undecidable and
+## Julia reports ambiguities; the per-eltype signatures dominate them.
+for ET in (TracedRInteger, TracedRFloat, TracedRComplex, TracedRReal, TracedRNumber)
+    @eval function Base.getindex(
+        a::AbstractArray{$ET{T},N}, index::Vararg{Union{Int,TracedRNumber{Int}},N}
+    ) where {T,N}
+        ancestor, idxs = TracedUtils.get_ancestor_and_indices(a, index...)
+        return getindex(ancestor, idxs...)
+    end
 end
 
 # This method is needed exclusively to resolve an ambiguity
@@ -88,76 +97,73 @@ function Base.getindex(
     return getindex(Reactant.promote_to(TracedRArray{T,1}, l), index)
 end
 
+## Scalar range indexing with traced indices (and traced-eltype ranges with
+## plain indices): one implementation per range family, with the Base-facing
+## `getindex` methods below as pure dispatch shims.
+function traced_range_getindex(r::AbstractRange, index)
+    return convert(
+        TracedRNumber{Reactant.unwrapped_eltype(r)},
+        first(r) + (index - oneunit(index)) * Base.step_hp(r),
+    )
+end
+function traced_range_getindex(r::Base.OneTo, index)
+    return convert(TracedRNumber{Reactant.unwrapped_eltype(r)}, index)
+end
+function traced_range_getindex(r::StepRangeLen, index)
+    # FIXME: this crashes for some reason
+    # u = convert(TracedRNumber{Reactant.unwrapped_eltype(r)}, index) - r.offset
+    # return convert(TracedRNumber{Reactant.unwrapped_eltype(r)}, r.ref + u * r.step)
+    return @allowscalar getindex(
+        Reactant.promote_to(TracedRArray{Reactant.unwrapped_eltype(r),1}, r), index
+    )
+end
+function traced_range_getindex(r::LinRange, index)
+    # `Base.lerpi` casts its result through the element type, which a traced
+    # interpolant cannot satisfy; interpolate with traced arithmetic instead.
+    t = (index - oneunit(index)) / r.lendiv
+    return convert(
+        TracedRNumber{Reactant.unwrapped_eltype(r)}, (1 - t) * r.start + t * r.stop
+    )
+end
+
 for (aT, iT) in (
     (AbstractRange, TracedRInteger),
     (AbstractRange{<:TracedRNumber}, Int),
     (AbstractRange{<:TracedRNumber}, TracedRInteger),
     (AbstractRange{<:TracedRNumber}, Union{Int,TracedRInteger{Int}}),
     (Base.IdentityUnitRange, TracedRInteger),
-    # 1.10 specific ambiguity fixes
-    (UnitRange{<:TracedRNumber}, Int),
-    # 1.10's `getindex(::UnitRange{<:Union{...}}, ::Integer)` overflow-safe
-    # grouping (range.jl) needs an exact-grouping disambiguator now that
-    # traced integers are `Integer`s
-    (
-        UnitRange{
-            <:Union{Bool,Int128,Int16,Int32,Int64,Int8,UInt128,UInt16,UInt32,UInt64,UInt8}
-        },
-        TracedRInteger,
-    ),
 )
-    @eval function Base.getindex(a::$aT, index::$iT)
-        return convert(
-            TracedRNumber{Reactant.unwrapped_eltype(a)},
-            first(a) + (index - oneunit(index)) * Base.step_hp(a),
-        )
-    end
+    @eval Base.getindex(r::$aT, index::$iT) = traced_range_getindex(r, index)
 end
 
-for (aT, iT) in (
-    (Base.OneTo, TracedRInteger),
-    (Base.OneTo{<:TracedRNumber}, Int),
-    (Base.OneTo{<:TracedRNumber}, TracedRInteger),
-    (Base.OneTo{<:TracedRNumber}, TracedRInteger{Int}),
-)
-    @eval function Base.getindex(a::$aT, index::$iT)
-        return convert(TracedRNumber{Reactant.unwrapped_eltype(a)}, index)
-    end
-end
-
-for (aT, iT) in (
-    (StepRangeLen, TracedRInteger),
-    (StepRangeLen{<:TracedRNumber}, TracedRInteger),
-    (StepRangeLen{<:TracedRNumber}, TracedRInteger{Int}),
-)
-    @eval function Base.getindex(r::$aT, index::$iT)
-        # FIXME: this crashes for some reason
-        # u = convert(TracedRNumber{Reactant.unwrapped_eltype(r)}, index) - r.offset
-        # return convert(TracedRNumber{Reactant.unwrapped_eltype(r)}, r.ref + u * r.step)
-        return @allowscalar getindex(
-            Reactant.promote_to(TracedRArray{Reactant.unwrapped_eltype(r),1}, r), index
-        )
-    end
-end
-
-for (aT, iT) in (
-    (LinRange, TracedRInteger),
-    (LinRange{<:TracedRNumber}, TracedRInteger),
-    (LinRange{<:TracedRNumber}, TracedRInteger{Int}),
-)
-    @eval function Base.getindex(r::$aT, index::$iT)
-        return convert(
-            TracedRNumber{Reactant.unwrapped_eltype(r)},
-            Base.lerpi(index - oneunit(index), r.lendiv, r.start, r.stop),
-        )
+## Julia 1.10's specialized range `getindex(::X{T}, ::Integer)` methods bind
+## the element type as a method typevar, and a disambiguator only dominates
+## them if it binds it the same way (`LinRange` and `StepRangeLen` share one
+## union-shaped method there). Cover each specialized range family against
+## each traced index shape.
+for R in (:(UnitRange{T}), :(Base.OneTo{T}), :(Union{LinRange{T},StepRangeLen{T}}))
+    @eval begin
+        Base.getindex(r::$R, index::TracedRInteger) where {T} =
+            traced_range_getindex(r, index)
+        function Base.getindex(r::$R, index::TracedRInteger) where {T<:TracedRNumber}
+            return traced_range_getindex(r, index)
+        end
+        function Base.getindex(r::$R, index::Int) where {T<:TracedRNumber}
+            return traced_range_getindex(r, index)
+        end
+        function Base.getindex(
+            r::$R, index::Union{Int,TracedRInteger{Int}}
+        ) where {T<:TracedRNumber}
+            return traced_range_getindex(r, index)
+        end
     end
 end
 
-# v1.10 specific ambiguity fixes
+# 1.10's overflow-safe `UnitRange` grouping needs its own exact-grouping cover
 function Base.getindex(
-    a::Union{LinRange{TracedRNumber{T}},StepRangeLen{TracedRNumber{T}}}, index::Int
-) where {T}
-    return getindex(Reactant.promote_to(TracedRArray{T,1}, a), index)
+    r::UnitRange{T}, index::TracedRInteger
+) where {T<:Union{Bool,Int128,Int16,Int32,Int64,Int8,UInt128,UInt16,UInt32,UInt64,UInt8}}
+    return traced_range_getindex(r, index)
 end
 
 function Base.getindex(a::AnyTracedRArray{T,N}, linear_indices) where {T,N}
diff --git a/src/stdlibs/LinearAlgebra.jl b/src/stdlibs/LinearAlgebra.jl
index 1c4b343763..f625a6ddf0 100644
--- a/src/stdlibs/LinearAlgebra.jl
+++ b/src/stdlibs/LinearAlgebra.jl
@@ -3,7 +3,12 @@ module TracedLinearAlgebra
 using ..MLIR: MLIR
 using ..Reactant: Reactant, Ops
 using ..Reactant:
-    TracedRArray, TracedRNumber, AnyTracedRArray, AnyTracedRMatrix, AnyTracedRVector
+    TracedRArray,
+    TracedRNumber,
+    TracedRInteger,
+    AnyTracedRArray,
+    AnyTracedRMatrix,
+    AnyTracedRVector
 using ..Reactant: call_with_reactant, unwrapped_eltype, promote_to
 using ReactantCore: ReactantCore, materialize_traced_array, @trace
 using Reactant_jll: Reactant_jll
@@ -714,8 +719,18 @@ for AT in (
     LinearAlgebra.UnitLowerTriangular{<:TracedRNumber},
     LinearAlgebra.UpperHessenberg{<:TracedRNumber},
 )
-    @eval function Base.getindex(A::$AT, i::Int, j::Int)
-        return getindex(materialize_traced_array(A), i, j)
+    @eval begin
+        function Base.getindex(A::$AT, i::Int, j::Int)
+            return getindex(materialize_traced_array(A), i, j)
+        end
+        # 1.10's structured-matrix `getindex` methods take `Integer`s, so
+        # their intersection with the traced-array methods includes traced
+        # integer indices; 1.12's take `Int`s, which the first cover matches.
+        function Base.getindex(
+            A::$AT, i::Union{Int,TracedRInteger{Int}}, j::Union{Int,TracedRInteger{Int}}
+        )
+            return getindex(materialize_traced_array(A), i, j)
+        end
     end
 end
 
diff --git a/test/core/ranges.jl b/test/core/ranges.jl
index 51d3f703fb..44994efd86 100644
--- a/test/core/ranges.jl
+++ b/test/core/ranges.jl
@@ -8,6 +8,13 @@ using Reactant, Test
     @test Array{Int64}(@jit(i:j)) == collect(5:10)
 end
 
+@testset "range getindex with traced index" begin
+    k = Reactant.ConcreteRNumber(3)
+    @test (@jit getindex(2:10, k)) == 4
+    @test (@jit getindex(LinRange(0.0, 1.0, 11), k)) ≈ 0.2
+    @test (@jit getindex(range(0.0; step=0.5, length=10), k)) ≈ 1.0
+end
+
 broadcast_over_range(a, kx, ky) = a .* (kx .^ 2 .+ ky' .^ 2)
 
 @testset "broadcast over ranges" begin

From 385e395768bd68c79f428f7093a1dfe6c76506dc Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 9 Jul 2026 16:27:35 +0200
Subject: [PATCH 26/30] Disambiguate Bool rounding of traced floats

Julia 1.10's `round`/`floor`/`ceil(::Type{Bool}, ::AbstractFloat)` are
ambiguous with the `::Type{T<:Integer}` traced-float methods; add the
exact `Bool` disambiguators.

Created by generative AI.
---
 src/TracedRNumber.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index f59570e5c5..fba06cd785 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -1215,6 +1215,11 @@ function Base.ceil(::Type{T}, x::TracedRFloat) where {T<:Integer}
     return trunc(T, ceil(x))
 end
 
+# disambiguate against 1.10's `round/floor/ceil(::Type{Bool}, ::AbstractFloat)`
+Base.round(::Type{Bool}, x::TracedRFloat) = trunc(Bool, round(x))
+Base.floor(::Type{Bool}, x::TracedRFloat) = trunc(Bool, floor(x))
+Base.ceil(::Type{Bool}, x::TracedRFloat) = trunc(Bool, ceil(x))
+
 # Concatenation. Numbers in Julia are handled in a much less generic fashion than arrays
 Base.vcat(x::TracedRNumber...) = Base.typed_vcat(Base.promote_eltypeof(x...), x...)
 function Base.typed_vcat(::Type{T}, x::TracedRNumber...) where {T}

From 80f197e5b7cdd134a70632c6c53251d79127af0b Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Thu, 9 Jul 2026 16:28:46 +0200
Subject: [PATCH 27/30] Widen `Missing`/`Nothing` promotion to the traced
 number kinds

The base branch keeps these `promote_rule` methods invariant because
loosening them there makes their specificity against Base's
`Missing`/`Nothing` rules undecidable. Here `TracedRNumber{T}` is a
union of the concrete kind types, so the invariant signatures match
nothing, and the general traced promotion rules are `Number`-bounded, so
the covariant forms no longer criss-cross Base's rules.

Created by generative AI.
---
 src/TracedRNumber.jl | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index fba06cd785..badc50f5d9 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -197,57 +197,57 @@ function Base.promote_rule(::Type{Nothing}, ::Type{<:TracedRNumber{S}}) where {S
     return Union{Nothing,TracedRNumber{S}}
 end
 
-function Base.promote_rule(::Type{TracedRNumber{T}}, ::Type{Nothing}) where {T}
+function Base.promote_rule(::Type{<:TracedRNumber{T}}, ::Type{Nothing}) where {T}
     return Union{Nothing,TracedRNumber{T}}
 end
 
-function Base.promote_rule(::Type{Missing}, ::Type{TracedRNumber{S}}) where {S}
+function Base.promote_rule(::Type{Missing}, ::Type{<:TracedRNumber{S}}) where {S}
     return Union{Missing,TracedRNumber{S}}
 end
 
-function Base.promote_rule(::Type{TracedRNumber{T}}, ::Type{Missing}) where {T}
+function Base.promote_rule(::Type{<:TracedRNumber{T}}, ::Type{Missing}) where {T}
     return Union{Missing,TracedRNumber{T}}
 end
 
 function Base.promote_rule(
-    ::Type{Union{Nothing,Missing}}, ::Type{TracedRNumber{S}}
+    ::Type{Union{Nothing,Missing}}, ::Type{<:TracedRNumber{S}}
 ) where {S}
     return Union{Nothing,Missing,TracedRNumber{S}}
 end
 
 function Base.promote_rule(
-    ::Type{TracedRNumber{T}}, ::Type{Union{Nothing,Missing}}
+    ::Type{<:TracedRNumber{T}}, ::Type{Union{Nothing,Missing}}
 ) where {T}
     return Union{Nothing,Missing,TracedRNumber{T}}
 end
 
 function Base.promote_rule(
-    T::Type{>:Union{Nothing,Missing}}, ::Type{TracedRNumber{S}}
+    T::Type{>:Union{Nothing,Missing}}, ::Type{<:TracedRNumber{S}}
 ) where {S}
     T2 = nonmissingtype(Base.nonnothingtype(promote_rule(T, S)))
     return Union{Nothing,Missing,TracedRNumber{T2}}
 end
 
 function Base.promote_rule(
-    ::Type{TracedRNumber{T}}, S::Type{>:Union{Nothing,Missing}}
+    ::Type{<:TracedRNumber{T}}, S::Type{>:Union{Nothing,Missing}}
 ) where {T}
     T2 = nonmissingtype(Base.nonnothingtype(promote_rule(T, S)))
     return Union{Nothing,Missing,TracedRNumber{T2}}
 end
 
-function Base.promote_rule(T::Type{>:Missing}, ::Type{TracedRNumber{S}}) where {S}
+function Base.promote_rule(T::Type{>:Missing}, ::Type{<:TracedRNumber{S}}) where {S}
     return Union{Missing,TracedRNumber{nonmissingtype(promote_type(S, T))}}
 end
 
-function Base.promote_rule(::Type{TracedRNumber{T}}, S::Type{>:Missing}) where {T}
+function Base.promote_rule(::Type{<:TracedRNumber{T}}, S::Type{>:Missing}) where {T}
     return Union{Missing,TracedRNumber{nonmissingtype(promote_type(T, S))}}
 end
 
-function Base.promote_rule(T::Type{>:Nothing}, ::Type{TracedRNumber{S}}) where {S}
+function Base.promote_rule(T::Type{>:Nothing}, ::Type{<:TracedRNumber{S}}) where {S}
     return Union{Nothing,TracedRNumber{Base.nonnothingtype(promote_type(S, T))}}
 end
 
-function Base.promote_rule(::Type{TracedRNumber{T}}, S::Type{>:Nothing}) where {T}
+function Base.promote_rule(::Type{<:TracedRNumber{T}}, S::Type{>:Nothing}) where {T}
     return Union{Nothing,TracedRNumber{Base.nonnothingtype(promote_type(T, S))}}
 end
 

From 2e65d3587f2bd016ed1c92acb95c4b1392a9d0d6 Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 11:29:52 +0200
Subject: [PATCH 28/30] Keep the export list unchanged

Optional refinement. `RNumber`/`RReal` were never exported, so do not
export the new kind names either; the manual refers to them qualified.
The previously exported names all still exist.

Created by generative AI.
---
 docs/src/api/api.md       |  6 +++---
 src/Reactant.jl           | 11 +----------
 test/core/type_numbers.jl |  6 ++++++
 3 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/docs/src/api/api.md b/docs/src/api/api.md
index 9674527197..988207526d 100644
--- a/docs/src/api/api.md
+++ b/docs/src/api/api.md
@@ -36,9 +36,9 @@ ConcreteRArray
 ConcreteRNumber
 Reactant.RNumber
 Reactant.RReal
-RInteger
-RFloat
-RComplex
+Reactant.RInteger
+Reactant.RFloat
+Reactant.RComplex
 Reactant.traced_number_type
 Reactant.pjrt_number_type
 Reactant.ifrt_number_type
diff --git a/src/Reactant.jl b/src/Reactant.jl
index 2ba6d0f790..5fff723ebb 100644
--- a/src/Reactant.jl
+++ b/src/Reactant.jl
@@ -300,21 +300,12 @@ using .Compiler:
     code_xla,
     traced_getfield,
     compile
-export RInteger,
-    RFloat,
-    RComplex,
-    ConcreteRArray,
+export ConcreteRArray,
     ConcreteRNumber,
     ConcretePJRTArray,
     ConcretePJRTNumber,
-    ConcretePJRTInteger,
-    ConcretePJRTFloat,
-    ConcretePJRTComplex,
     ConcreteIFRTArray,
     ConcreteIFRTNumber,
-    ConcreteIFRTInteger,
-    ConcreteIFRTFloat,
-    ConcreteIFRTComplex,
     @compile,
     @code_hlo,
     @code_mhlo,
diff --git a/test/core/type_numbers.jl b/test/core/type_numbers.jl
index 858766bb76..bca705ee0e 100644
--- a/test/core/type_numbers.jl
+++ b/test/core/type_numbers.jl
@@ -289,6 +289,12 @@ end
         TracedRInteger,
         TracedRFloat,
         TracedRComplex,
+        ConcretePJRTInteger,
+        ConcretePJRTFloat,
+        ConcretePJRTComplex,
+        ConcreteIFRTInteger,
+        ConcreteIFRTFloat,
+        ConcreteIFRTComplex,
         traced_number_type
 
     @testset "abstract kinds" begin

From 371d59e04d207fd26900dcdff60e718d3ff9113e Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 11:34:37 +0200
Subject: [PATCH 29/30] Share the Base disambiguation-shape tables

Optional refinement. The rounding-mode method shapes Base specializes
`div` on were listed in two places; name them once. The overflow-safe
integer grouping is `Base.OverflowSafe`, whose lifecycle is tied to the
very `UnitRange` method the cover disambiguates against.

Created by generative AI.
---
 src/ConcreteRArray.jl | 12 ------------
 src/Indexing.jl       |  4 +---
 src/TracedRNumber.jl  | 13 +------------
 src/utils.jl          | 16 ++++++++++++++++
 4 files changed, 18 insertions(+), 27 deletions(-)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index eef3db099e..62a9ea7ad7 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -285,18 +285,6 @@ end
 
 # Base specializes `div` on single rounding modes and on `Rational` arguments;
 # the corresponding methods below only disambiguate against those.
-const BASE_SPECIFIC_ROUNDING_MODES = (
-    RoundingMode{:FromZero},
-    RoundingMode{:Nearest},
-    RoundingMode{:NearestTiesAway},
-    RoundingMode{:NearestTiesUp},
-    RoundingMode{:Up},
-    RoundingMode{:Down},
-    # Base also groups the nearest modes into a single method
-    Union{
-        RoundingMode{:Nearest},RoundingMode{:NearestTiesAway},RoundingMode{:NearestTiesUp}
-    },
-)
 
 for T1 in CONCRETE_NUMBER_KINDS
     @eval begin
diff --git a/src/Indexing.jl b/src/Indexing.jl
index e314572b2a..4eaca1c324 100644
--- a/src/Indexing.jl
+++ b/src/Indexing.jl
@@ -160,9 +160,7 @@ for R in (:(UnitRange{T}), :(Base.OneTo{T}), :(Union{LinRange{T},StepRangeLen{T}
 end
 
 # 1.10's overflow-safe `UnitRange` grouping needs its own exact-grouping cover
-function Base.getindex(
-    r::UnitRange{T}, index::TracedRInteger
-) where {T<:Union{Bool,Int128,Int16,Int32,Int64,Int8,UInt128,UInt16,UInt32,UInt64,UInt8}}
+function Base.getindex(r::UnitRange{T}, index::TracedRInteger) where {T<:Base.OverflowSafe}
     return traced_range_getindex(r, index)
 end
 
diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index badc50f5d9..8ae6adb5e7 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -557,18 +557,7 @@ Base.div(::Missing, @nospecialize(rhs::TracedRInteger), ::RoundingMode) = missin
 
 # Base has `div` methods specialized on single rounding modes (and on
 # `Rational` arguments); the following methods only disambiguate against them.
-for RM in (
-    RoundingMode{:FromZero},
-    RoundingMode{:Nearest},
-    RoundingMode{:NearestTiesAway},
-    RoundingMode{:NearestTiesUp},
-    RoundingMode{:Up},
-    RoundingMode{:Down},
-    # Base also groups the nearest modes into a single method
-    Union{
-        RoundingMode{:Nearest},RoundingMode{:NearestTiesAway},RoundingMode{:NearestTiesUp}
-    },
-)
+for RM in Reactant.BASE_SPECIFIC_ROUNDING_MODES
     @eval begin
         function Base.div(
             @nospecialize(lhs::TracedRInteger{T1}),
diff --git a/src/utils.jl b/src/utils.jl
index e293030cc5..da83f2c639 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -1204,3 +1204,19 @@ nmantissa(::Type{Float64}) = 52
 
 @inline _unwrap_val(x) = x
 @inline _unwrap_val(::Val{val}) where {val} = _unwrap_val(val)
+
+# Base specializes `div` (and friends) on these single-rounding-mode method
+# shapes; disambiguators against them must enumerate the same shapes.
+const BASE_SPECIFIC_ROUNDING_MODES = (
+    RoundingMode{:FromZero},
+    RoundingMode{:Nearest},
+    RoundingMode{:NearestTiesAway},
+    RoundingMode{:NearestTiesUp},
+    RoundingMode{:Up},
+    RoundingMode{:Down},
+    # Base also groups the nearest modes into a single method
+    Union{
+        RoundingMode{:Nearest},RoundingMode{:NearestTiesAway},RoundingMode{:NearestTiesUp}
+    },
+)
+

From 16303a1ebe22db9ebb2df84cf01cf06450e0067c Mon Sep 17 00:00:00 2001
From: Oliver Schulz <oschulz@mpp.mpg.de>
Date: Fri, 10 Jul 2026 11:45:53 +0200
Subject: [PATCH 30/30] Generate the number-forwarding grids from one iterator

Optional refinement. The concrete-number methods and the kernel-side
CUDA number methods forwarded the same operator grid (kind x kind,
`Number`, and each Base-specialized number type) in two hand-maintained
nested loops; `foreach_number_kind_forwarding` now walks the grid and
each site only states its operators and how to unwrap. The generated
method set is unchanged (verified signature-for-signature).

Created by generative AI.
---
 ext/ReactantCUDAExt.jl | 63 +++++++++++++++---------------------------
 src/ConcreteRArray.jl  | 48 +++++++++++---------------------
 src/utils.jl           | 44 ++++++++++++++++++++++++++++-
 3 files changed, 81 insertions(+), 74 deletions(-)

diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl
index bc27680e74..a0cb9f5480 100644
--- a/ext/ReactantCUDAExt.jl
+++ b/ext/ReactantCUDAExt.jl
@@ -122,47 +122,28 @@ Base.convert(::Type{T}, RN::CuTracedRNumber) where {T<:Number} = convert(T, geti
 # would be ambiguous with Base's methods on `Real`/`Integer`/`AbstractFloat`.
 const CU_TRACED_NUMBER_KINDS = (CuTracedRInteger, CuTracedRFloat, CuTracedRComplex)
 
-for jlop in (
-    :(Base.min),
-    :(Base.mod),
-    :(Base.max),
-    :(Base.:+),
-    :(Base.:-),
-    :(Base.:*),
-    :(Base.:/),
-    :(Base.:^),
-    :(Base.rem),
-    :(Base.isless),
-    :(Base.:(==)),
-    :(Base.:(!=)),
-)
-    for K1 in CU_TRACED_NUMBER_KINDS
-        @eval begin
-            @inline $jlop(a::$K1, b::Number) = $jlop(a[], b)
-            @inline $jlop(a::Number, b::$K1) = $jlop(a, b[])
-        end
-        for K2 in CU_TRACED_NUMBER_KINDS
-            @eval @inline $jlop(a::$K1, b::$K2) = $jlop(a[], b[])
-        end
-        for X in (
-            :Real,
-            :Integer,
-            :Bool,
-            :AbstractFloat,
-            :Complex,
-            :(Complex{Bool}),
-            :Rational,
-            :BigInt,
-            :BigFloat,
-            :AbstractIrrational,
-        )
-            jlop == :(Base.:^) && X === :Integer && continue
-            @eval begin
-                @inline $jlop(a::$X, b::$K1) = $jlop(a, b[])
-                @inline $jlop(a::$K1, b::$X) = $jlop(a[], b)
-            end
-        end
-    end
+Reactant.foreach_number_kind_forwarding(
+    (
+        :(Base.min),
+        :(Base.mod),
+        :(Base.max),
+        :(Base.:+),
+        :(Base.:-),
+        :(Base.:*),
+        :(Base.:/),
+        :(Base.:^),
+        :(Base.rem),
+        :(Base.isless),
+        :(Base.:(==)),
+        :(Base.:(!=)),
+    ),
+    CU_TRACED_NUMBER_KINDS;
+    # `^` with a plain integer exponent has its own constant-propagating method
+    skip=(jlop, X) -> jlop == :(Base.:^) && X === :Integer,
+) do jlop, TX, TY, unwrap_a, unwrap_b
+    aa = unwrap_a ? :(a[]) : :a
+    ab = unwrap_b ? :(b[]) : :b
+    @eval @inline $(jlop)(a::$(TX), b::$(TY)) = $(jlop)($(aa), $(ab))
 end
 
 @inline Base.ifelse(cond::Bool, a, b::CuTracedRNumber) = ifelse(cond, a, b[])
diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 62a9ea7ad7..341a4b24a4 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -226,7 +226,7 @@ const CONCRETE_NUMBER_KINDS = (
 )
 const TRACED_NUMBER_KINDS = (TracedRInteger, TracedRFloat, TracedRComplex)
 
-for jlop in (
+const FORWARDED_NUMBER_OPS = (
     :(Base.isless),
     :(Base.:<),
     :(Base.:<=),
@@ -239,39 +239,23 @@ for jlop in (
     :(Base.:^),
     :(Base.:(==)),
 )
-    for T1 in CONCRETE_NUMBER_KINDS
-        @eval begin
-            $(jlop)(x::$(T1), y::Number) = $(jlop)(to_number(x), y)
-            $(jlop)(x::Number, y::$(T1)) = $(jlop)(x, to_number(y))
-        end
-        for T2 in CONCRETE_NUMBER_KINDS
-            @eval $(jlop)(x::$(T1), y::$(T2)) = $(jlop)(to_number(x), to_number(y))
-        end
-        for X in (
-            :Real,
-            :Integer,
-            :Bool,
-            :Rational,
-            :BigInt,
-            :BigFloat,
-            :AbstractIrrational,
-            :AbstractFloat,
-            :Complex,
-            :(Complex{Bool}),
-        )
-            @eval begin
-                $(jlop)(x::$(T1), y::$(X)) = $(jlop)(to_number(x), y)
-                $(jlop)(x::$(X), y::$(T1)) = $(jlop)(x, to_number(y))
-            end
-        end
-        for T2 in TRACED_NUMBER_KINDS
-            @eval begin
-                $(jlop)(x::$(T1), y::$(T2)) = throw(MethodError($(jlop), (x, y)))
-                $(jlop)(x::$(T2), y::$(T1)) = throw(MethodError($(jlop), (x, y)))
-            end
-        end
+
+foreach_number_kind_forwarding(
+    FORWARDED_NUMBER_OPS, CONCRETE_NUMBER_KINDS
+) do jlop, TX, TY, unwrap_x, unwrap_y
+    ax = unwrap_x ? :(to_number(x)) : :x
+    ay = unwrap_y ? :(to_number(y)) : :y
+    @eval $(jlop)(x::$(TX), y::$(TY)) = $(jlop)($(ax), $(ay))
+end
+
+for jlop in FORWARDED_NUMBER_OPS, T1 in CONCRETE_NUMBER_KINDS, T2 in TRACED_NUMBER_KINDS
+    @eval begin
+        $(jlop)(x::$(T1), y::$(T2)) = throw(MethodError($(jlop), (x, y)))
+        $(jlop)(x::$(T2), y::$(T1)) = throw(MethodError($(jlop), (x, y)))
     end
+end
 
+for jlop in FORWARDED_NUMBER_OPS
     T = AbstractConcreteArray{<:Number,0}
     @eval begin
         $(jlop)(x::$(T), y::$(T)) = $(jlop)(to_number(x), to_number(y))
diff --git a/src/utils.jl b/src/utils.jl
index da83f2c639..b68745baec 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -1205,6 +1205,49 @@ nmantissa(::Type{Float64}) = 52
 @inline _unwrap_val(x) = x
 @inline _unwrap_val(::Val{val}) where {val} = _unwrap_val(val)
 
+# Base and the stdlibs specialize the binary number operators on many types
+# (`Real`, `Rational`, `BigFloat`, ...). A wrapper-number kind that subtypes
+# `Integer`/`AbstractFloat` must forward such operators per kind and per
+# specialized argument type: a single method on the kind union would be
+# ambiguous with the specialized methods.
+const BASE_SPECIFIC_NUMBER_TYPES = (
+    :Real,
+    :Integer,
+    :Bool,
+    :Rational,
+    :BigInt,
+    :BigFloat,
+    :AbstractIrrational,
+    :AbstractFloat,
+    :Complex,
+    :(Complex{Bool}),
+)
+
+"""
+    foreach_number_kind_forwarding(f, ops, kinds; skip=(op, X) -> false)
+
+Call `f(op, TX, TY, unwrap_x, unwrap_y)` once per method signature required to
+forward the binary operators `ops` for the wrapper-number `kinds`: kind-kind
+pairs, `Number` pairs, and pairs with each Base-specialized number type (see
+`BASE_SPECIFIC_NUMBER_TYPES`; `skip(op, X)` exempts an operator from one of
+them). The `unwrap` flags mark which arguments are wrapper numbers.
+"""
+function foreach_number_kind_forwarding(f, ops, kinds; skip=(op, X) -> false)
+    for op in ops, K1 in kinds
+        f(op, K1, :Number, true, false)
+        f(op, :Number, K1, false, true)
+        for K2 in kinds
+            f(op, K1, K2, true, true)
+        end
+        for X in BASE_SPECIFIC_NUMBER_TYPES
+            skip(op, X) && continue
+            f(op, K1, X, true, false)
+            f(op, X, K1, false, true)
+        end
+    end
+    return nothing
+end
+
 # Base specializes `div` (and friends) on these single-rounding-mode method
 # shapes; disambiguators against them must enumerate the same shapes.
 const BASE_SPECIFIC_ROUNDING_MODES = (
@@ -1219,4 +1262,3 @@ const BASE_SPECIFIC_ROUNDING_MODES = (
         RoundingMode{:Nearest},RoundingMode{:NearestTiesAway},RoundingMode{:NearestTiesUp}
     },
 )
-