Add ShapeDtypeStruct for compilation without concrete array allocation

docs/src/api/api.md

@@ -32,6 +32,8 @@ Reactant.to_rarray
 ```@docs
 ConcreteRArray
 ConcreteRNumber
+RArraySpec
+RNumberSpec
 ```
 
 ## Inspect Generated HLO

src/Compiler.jl

@@ -3947,7 +3947,12 @@ function __resolve_device_and_client(client, seen_args, linear_args, is_sharded)
     if length(linear_args) > 0
         devices_list = []
         for (k, v) in seen_args
-            !(v isa TracedRArray || v isa TracedRNumber) && continue
+            if (
+                !(v isa TracedRArray || v isa TracedRNumber) ||
+                (k isa Reactant.RArraySpec || k isa Reactant.RNumberSpec)
+            )
+                continue
+            end
             buffer = k.data isa Tuple ? only(k.data) : k.data
             push!(devices_list, XLA.device(buffer))
         end

src/ConcreteRArray.jl

@@ -286,6 +286,23 @@ function Base.show(io::IO, X::Union{ConcretePJRTScalar,ConcreteIFRTScalar})
     return nothing
 end
 
+function Base.showarg(io::IO, ::RArraySpec{T,N}, toplevel) where {T,N}
+    toplevel || print(io, "::")
+    print(io, "RArraySpec{$T,$N}")
+    # TODO: Add sharding info
+    return nothing
+end
+
+function Base.print_array(io::IO, ::RArraySpec)
+    print(io, "<RArraySpec has no data>")
+    return nothing
+end
+
+function Base.show(io::IO, X::RArraySpec)
+    print(io, "$(typeof(X))()")
+    return nothing
+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))>")

src/Reactant.jl

@@ -124,6 +124,9 @@ unwrapped_eltype(::TracedRNumber{T}) where {T} = T
 unwrapped_eltype(::Type{<:AbstractArray{T,N}}) where {T,N} = unwrapped_eltype(T)
 unwrapped_eltype(::AbstractArray{T,N}) where {T,N} = unwrapped_eltype(T)
 
+unwrapped_eltype(::Type{RArraySpec{T,N}}) where {T,N} = T
+unwrapped_eltype(::RArraySpec{T,N}) where {T,N} = T
+
 include("Ops.jl")
 Base.push!(no_rewrite_ancestor_modules, Ops)
 
@@ -289,6 +292,8 @@ export ConcreteRArray,
     ConcretePJRTNumber,
     ConcreteIFRTArray,
     ConcreteIFRTNumber,
+    RArraySpec,
+    RNumberSpec,
     @compile,
     @code_hlo,
     @code_mhlo,