Skip to content

Groundwork for number type separation #3034

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
oschulz wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Groundwork for number type separation #3034

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5bbfb25
Widen traced-type signatures to accept covariant element types
oschulz Jul 8, 2026
25f12ff
Make `similar` of a concrete array with a traced eltype return a trac…
oschulz Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ext/ReactantAbstractFFTsExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions ext/ReactantCUDAExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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},
Expand Down
6 changes: 3 additions & 3 deletions ext/ReactantSparseArraysExt/Errors.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
34 changes: 17 additions & 17 deletions ext/ReactantSparseArraysExt/ReadOnly.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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...)
Expand All @@ -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...)
Expand Down
2 changes: 1 addition & 1 deletion lib/ReactantCore/src/ReactantCore.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
22 changes: 18 additions & 4 deletions src/ConcreteRArray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -725,8 +739,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
Expand Down Expand Up @@ -771,7 +785,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
Expand Down Expand Up @@ -805,7 +819,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
Expand Down
14 changes: 7 additions & 7 deletions src/Indexing.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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}}...)
Expand All @@ -188,16 +188,16 @@ 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)
end
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
Expand Down Expand Up @@ -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)))))
Expand Down
22 changes: 11 additions & 11 deletions src/Ops.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down Expand Up @@ -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}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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}
Expand All @@ -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}
Expand Down Expand Up @@ -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__),
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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},
Expand Down
Loading
Loading