Fix typedef and undef ctors for ConcreteRArray and ConcreteRNumber (#1571)

* Fix typedef of ConcreteRNumber

* Fix typedef and ctors of ConcreteRArray and add tests

* Add tests for ConcreteRNumber

* Remove deprecated ConcreteRArray undef ctor with T as argument

* Don't restrict ConcreteRArray undef ctor to eltype Number

* Remove type parameters in ConcreteRArray and ConcreteRNumber typedefs

Author: oschulz

ext/ReactantKernelAbstractionsExt.jl

@@ -26,7 +26,7 @@ function Base.getproperty(x::ReactantBackend, sym::Symbol)
 end
 
 function KA.allocate(::ReactantBackend, ::Type{T}, dims::Tuple) where {T}
-    return ConcreteRArray(undef, T, dims)
+    return ConcreteRArray{T}(undef, dims)
 end
 
 function KA.zeros(b::ReactantBackend, ::Type{T}, dims::Tuple) where {T}

src/ConcreteRArray.jl

@@ -378,8 +378,8 @@ end
     device::Union{Nothing,XLA.PJRT.Device}=nothing,
     sharding::Sharding.AbstractSharding=Sharding.NoSharding(),
 ) where {S}
-    return ConcretePJRTArray(
-        undef, S, dims; client=client, idx=idx, device=device, sharding=sharding
+    return ConcretePJRTArray{S}(
+        undef, dims; client=client, idx=idx, device=device, sharding=sharding
     )
 end
 
@@ -410,7 +410,7 @@ function Base.similar(a::ConcreteIFRTArray{T}, ::Type{S}=T, dims::Dims=size(a))
 end
 Base.similar(a::ConcreteIFRTArray, dims::Dims) = similar(a, eltype(a), dims)
 function Base.similar(::Type{ConcreteIFRTArray{T}}, dims) where {T}
-    return ConcreteIFRTArray(undef, T, dims)
+    return ConcreteIFRTArray{T}(undef, dims)
 end
 
 # Broadcasting interface

src/Types.jl

@@ -410,14 +410,16 @@ Base.@deprecate_binding ConcreteRNG ReactantRNG
 Base.@deprecate_binding TracedRNG ReactantRNG
 
 """
-    ConcreteRArray(
-        undef, ::Type{T}, shape::Dims;
+    ConcreteRArray{T}(
+        undef, shape::Dims;
         client::Union{Nothing,XLA.AbstractClient} = nothing,
         device::Union{Nothing,XLA.AbstractDevice} = nothing,
         sharding::Sharding.AbstractSharding = Sharding.NoSharding(),
     )
 
-    ConcretePJRTArray(data::Array, kwargs...)
+    ConcretePJRTArray{T}(undef, shape::Integer...; kwargs...)
+
+    ConcretePJRTArray(data::Array; kwargs...)
 
 Allocate an uninitialized `ConcreteRArray` of element type `T` and size
 `shape` or convert an `Array` to a `ConcreteRArray`.
@@ -434,6 +436,9 @@ elseif XLA.REACTANT_XLA_RUNTIME == "IFRT"
     ConcreteIFRTArray
 end
 
+@inline ConcreteRArray{T}(::UndefInitializer, shape::Integer...; kwargs...) where {T} =
+    ConcreteRArray{T}(undef, Dims(shape); kwargs...)
+
 """
     ConcreteRNumber(
         x::Number;
@@ -460,16 +465,13 @@ end
 
 ## Other Aliases based on the set preferences
 @static if XLA.REACTANT_XLA_RUNTIME == "PJRT"
-    const ConcreteRNumber = ConcretePJRTNumber
     const AnyConcreteRArray = AnyConcretePJRTArray
 elseif XLA.REACTANT_XLA_RUNTIME == "IFRT"
-    const ConcreteRNumber = ConcreteIFRTNumber
     const AnyConcreteRArray = AnyConcreteIFRTArray
 end
 
-function ConcretePJRTArray(
+function ConcretePJRTArray{T}(
     ::UndefInitializer,
-    ::Type{T},
     shape::Dims;
     client::Union{Nothing,XLA.AbstractClient}=nothing,
     idx::Union{Int,Nothing}=nothing,
@@ -483,9 +485,8 @@ function ConcretePJRTArray(
     return ConcretePJRTArray{T,N,nsharded,typeof(shardinfo)}(sharded_data, shape, shardinfo)
 end
 
-function ConcreteIFRTArray(
+function ConcreteIFRTArray{T}(
     ::UndefInitializer,
-    ::Type{T},
     shape::Dims;
     client::Union{Nothing,XLA.AbstractClient}=nothing,
     idx::Union{Int,Nothing}=nothing,

test/layout.jl

@@ -2,6 +2,34 @@ using Reactant
 using Test
 
 @testset "Layout" begin
+    client = Reactant.XLA.default_backend()
+    device = Reactant.XLA.default_device()
+    sharding = Sharding.NoSharding()
+    idx = 0
+
+    @test ConcreteRArray{Float32}(undef, (100, 10)) isa ConcreteRArray{Float32,2}
+
+    @test ConcreteRArray{Float32}(
+        undef, (100, 10); client=client, idx=idx, device=device
+    ) isa ConcreteRArray{Float32,2}
+
+    @test ConcreteRArray{Float32}(
+        undef, Int32(100), Int16(10); client=client, idx=idx, device=device
+    ) isa ConcreteRArray{Float32,2}
+
+    @test ConcreteRNumber(Float32(4.2)) isa ConcreteRNumber{Float32}
+
+    @test ConcreteRNumber(Float16(4.2); client=client, idx=idx, device=device) isa
+        ConcreteRNumber{Float16}
+
+    @test ConcreteRNumber{Float32}(Float32(4.2); client=client, idx=idx, device=device) isa
+        ConcreteRNumber{Float32}
+
+    @test ConcreteRNumber{Float16}(Float32(4.2)) isa ConcreteRNumber{Float16}
+
+    @test ConcreteRNumber{Float32}(Float16(4.2); client=client, idx=idx, device=device) isa
+        ConcreteRNumber{Float32}
+
     x = reshape([1.0, 2.0, 3.0, 4.0], (2, 2))
 
     y = Reactant.to_rarray(x)