From 74d57df258a96821d06c9813d4dfffde741b5106 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 12 Feb 2026 15:30:42 +0000
Subject: [PATCH 1/8] Initial plan


From ba6c67a9f11acd5efd51e2307cee91f501968a8a Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 12 Feb 2026 15:36:23 +0000
Subject: [PATCH 2/8] Add ShapeDtypeStruct type for compile without concrete
 arrays

Co-authored-by: avik-pal <30564094+avik-pal@users.noreply.github.com>
---
 src/Reactant.jl      |  1 +
 src/Tracing.jl       | 24 ++++++++++++++++
 src/Types.jl         | 49 +++++++++++++++++++++++++++++++++
 test/core/compile.jl | 65 ++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 139 insertions(+)

diff --git a/src/Reactant.jl b/src/Reactant.jl
index 6d31c41c04..11c1323149 100644
--- a/src/Reactant.jl
+++ b/src/Reactant.jl
@@ -289,6 +289,7 @@ export ConcreteRArray,
     ConcretePJRTNumber,
     ConcreteIFRTArray,
     ConcreteIFRTNumber,
+    ShapeDtypeStruct,
     @compile,
     @code_hlo,
     @code_mhlo,
diff --git a/src/Tracing.jl b/src/Tracing.jl
index ef96f2412a..2536ceae46 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -1407,6 +1407,30 @@ Base.@nospecializeinfer function make_tracer(
     return res
 end
 
+Base.@nospecializeinfer function make_tracer(
+    seen,
+    @nospecialize(prev::ShapeDtypeStruct{T,N}),
+    @nospecialize(path),
+    mode;
+    kwargs...,
+) where {T,N}
+    if mode == TracedToTypes
+        throw("Cannot have ShapeDtypeStruct as function call argument.")
+    end
+    if mode == ArrayToConcrete
+        throw("Cannot convert ShapeDtypeStruct to ConcreteRArray. ShapeDtypeStruct is only for compilation signatures.")
+    end
+    # ShapeDtypeStruct behaves like ConcreteToTraced mode - creates a TracedRArray without data
+    # Accept both ConcreteToTraced and TracedSetPath modes
+    if mode != ConcreteToTraced && mode != TracedSetPath
+        throw("ShapeDtypeStruct can only be used with ConcreteToTraced or TracedSetPath mode, got $mode")
+    end
+    haskey(seen, prev) && return seen[prev]::TracedRArray{T,N}
+    res = TracedRArray{T,N}((path,), nothing, size(prev))
+    seen[prev] = res
+    return res
+end
+
 Base.@nospecializeinfer function make_tracer(
     seen,
     prev::ConcretePJRTNumber{T},
diff --git a/src/Types.jl b/src/Types.jl
index 409a69d7cc..a2a3e0ca93 100644
--- a/src/Types.jl
+++ b/src/Types.jl
@@ -135,6 +135,55 @@ const AnyTracedRVector{T} = AnyTracedRArray{T,1}
 const AnyTracedRMatrix{T} = AnyTracedRArray{T,2}
 const AnyTracedRVecOrMat{T} = Union{AnyTracedRVector{T},AnyTracedRMatrix{T}}
 
+## ShapeDtypeStruct
+"""
+    ShapeDtypeStruct{T,N}(shape::NTuple{N,Int})
+    ShapeDtypeStruct(shape::NTuple{N,Int}, dtype::Type{T}) where {T,N}
+
+Lightweight structure that specifies the shape and element type (dtype) of an array
+without allocating the actual array data. Similar to JAX's `ShapeDtypeStruct`.
+
+This is useful for compiling functions without constructing the full `ConcreteRArray`,
+which can save memory and improve compilation performance.
+
+# Examples
+```julia
+# Specify shape and dtype for a 2D array
+spec = Reactant.ShapeDtypeStruct((10, 20), Float32)
+
+# Compile a function using just the spec
+f(x) = sum(x)
+compiled_f = Reactant.compile(f, (spec,))
+
+# Execute with actual data
+x = Reactant.ConcreteRArray(rand(Float32, 10, 20))
+result = compiled_f(x)
+```
+
+See also: [`compile`](@ref), [`ConcreteRArray`](@ref)
+"""
+struct ShapeDtypeStruct{T,N}
+    shape::NTuple{N,Int}
+
+    function ShapeDtypeStruct{T,N}(shape::NTuple{N,Int}) where {T,N}
+        return new{T,N}(shape)
+    end
+end
+
+function ShapeDtypeStruct(shape::NTuple{N,Int}, dtype::Type{T}) where {T,N}
+    return ShapeDtypeStruct{T,N}(shape)
+end
+
+function ShapeDtypeStruct(shape::Tuple{Vararg{Integer}}, dtype::Type{T}) where {T}
+    return ShapeDtypeStruct(map(Int, shape), dtype)
+end
+
+Base.size(x::ShapeDtypeStruct) = x.shape
+Base.ndims(::ShapeDtypeStruct{T,N}) where {T,N} = N
+Base.eltype(::ShapeDtypeStruct{T}) where {T} = T
+
+@leaf ShapeDtypeStruct
+
 # Concrete Types
 ## ConcretePJRTNumber
 mutable struct ConcretePJRTNumber{T,D} <: AbstractConcreteNumber{T}
diff --git a/test/core/compile.jl b/test/core/compile.jl
index a3a07c4300..1993ff96b8 100644
--- a/test/core/compile.jl
+++ b/test/core/compile.jl
@@ -640,3 +640,68 @@ end
     @test Array(y[:Mhalo]) ≈ [1.0f0, 2.0f0]
     @test Array(y[:x]) ≈ [2.0f0, 3.0f0]
 end
+
+@testset "ShapeDtypeStruct compilation" begin
+    @testset "Basic compilation with ShapeDtypeStruct" begin
+        # Define a simple function
+        f(x) = sum(x)
+        
+        # Compile using ShapeDtypeStruct instead of ConcreteRArray
+        spec = Reactant.ShapeDtypeStruct((10, 20), Float32)
+        compiled_f = Reactant.compile(f, (spec,))
+        
+        # Execute with actual data
+        x = Reactant.ConcreteRArray(rand(Float32, 10, 20))
+        result = compiled_f(x)
+        
+        @test result isa Reactant.ConcreteRNumber{Float32}
+        @test result ≈ sum(Array(x))
+    end
+    
+    @testset "Multiple arguments with ShapeDtypeStruct" begin
+        f(x, y) = x .+ y
+        
+        spec1 = Reactant.ShapeDtypeStruct((5, 5), Float64)
+        spec2 = Reactant.ShapeDtypeStruct((5, 5), Float64)
+        compiled_f = Reactant.compile(f, (spec1, spec2))
+        
+        x = Reactant.ConcreteRArray(rand(Float64, 5, 5))
+        y = Reactant.ConcreteRArray(rand(Float64, 5, 5))
+        result = compiled_f(x, y)
+        
+        @test result isa Reactant.ConcreteRArray{Float64,2}
+        @test result ≈ Array(x) .+ Array(y)
+    end
+    
+    @testset "ShapeDtypeStruct with different dtypes" begin
+        f(x) = sin.(x)
+        
+        for dtype in [Float32, Float64]
+            spec = Reactant.ShapeDtypeStruct((10,), dtype)
+            compiled_f = Reactant.compile(f, (spec,))
+            
+            x = Reactant.ConcreteRArray(rand(dtype, 10))
+            result = compiled_f(x)
+            
+            @test result isa Reactant.ConcreteRArray{dtype,1}
+            @test result ≈ sin.(Array(x))
+        end
+    end
+    
+    @testset "ShapeDtypeStruct constructor variations" begin
+        # Test different constructor forms
+        spec1 = Reactant.ShapeDtypeStruct{Float32,2}((3, 4))
+        @test size(spec1) == (3, 4)
+        @test eltype(spec1) == Float32
+        @test ndims(spec1) == 2
+        
+        spec2 = Reactant.ShapeDtypeStruct((3, 4), Float32)
+        @test size(spec2) == (3, 4)
+        @test eltype(spec2) == Float32
+        
+        # Test with integer tuple (not Int tuple)
+        spec3 = Reactant.ShapeDtypeStruct((3, 4), Float64)
+        @test size(spec3) == (3, 4)
+        @test eltype(spec3) == Float64
+    end
+end

From 7544d8be380345649bc0aaac9c7f8487204e4533 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 12 Feb 2026 15:40:28 +0000
Subject: [PATCH 3/8] Add unwrapped_eltype support for ShapeDtypeStruct

Co-authored-by: avik-pal <30564094+avik-pal@users.noreply.github.com>
---
 src/Reactant.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/Reactant.jl b/src/Reactant.jl
index 11c1323149..219eb60288 100644
--- a/src/Reactant.jl
+++ b/src/Reactant.jl
@@ -124,6 +124,10 @@ 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)
 
+# For ShapeDtypeStruct
+unwrapped_eltype(::Type{ShapeDtypeStruct{T,N}}) where {T,N} = T
+unwrapped_eltype(::ShapeDtypeStruct{T,N}) where {T,N} = T
+
 include("Ops.jl")
 Base.push!(no_rewrite_ancestor_modules, Ops)
 

From 2c8f6603e31825e1e92d05a3758184505d4a58fc Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 12 Feb 2026 15:41:22 +0000
Subject: [PATCH 4/8] Add documentation and examples for ShapeDtypeStruct

Co-authored-by: avik-pal <30564094+avik-pal@users.noreply.github.com>
---
 docs/src/api/api.md            |  1 +
 examples/shapedtype_example.jl | 94 ++++++++++++++++++++++++++++++++++
 2 files changed, 95 insertions(+)
 create mode 100644 examples/shapedtype_example.jl

diff --git a/docs/src/api/api.md b/docs/src/api/api.md
index 2c0bf8eddf..edef9a44fa 100644
--- a/docs/src/api/api.md
+++ b/docs/src/api/api.md
@@ -32,6 +32,7 @@ Reactant.to_rarray
 ```@docs
 ConcreteRArray
 ConcreteRNumber
+ShapeDtypeStruct
 ```
 
 ## Inspect Generated HLO
diff --git a/examples/shapedtype_example.jl b/examples/shapedtype_example.jl
new file mode 100644
index 0000000000..24d4a11b89
--- /dev/null
+++ b/examples/shapedtype_example.jl
@@ -0,0 +1,94 @@
+# Example: Using ShapeDtypeStruct for compilation without concrete arrays
+#
+# This example demonstrates how to use Reactant.ShapeDtypeStruct to compile
+# functions without having to construct full ConcreteRArray instances with
+# actual data. This is useful for:
+# 1. Faster compilation when you only need shape/dtype information
+# 2. Memory efficiency when working with large arrays
+# 3. Similar workflow to JAX's ShapeDtypeStruct
+
+using Reactant
+
+# Example 1: Basic usage with a simple function
+println("Example 1: Basic compilation with ShapeDtypeStruct")
+println("=" ^ 60)
+
+# Define a simple function that sums an array
+f_sum(x) = sum(x)
+
+# Instead of creating a full ConcreteRArray:
+# x = Reactant.ConcreteRArray(rand(Float32, 10, 20))  # This allocates memory!
+
+# Use ShapeDtypeStruct to specify only shape and dtype:
+spec = Reactant.ShapeDtypeStruct((10, 20), Float32)
+println("Created ShapeDtypeStruct: ", spec)
+println("  Shape: ", size(spec))
+println("  Element type: ", eltype(spec))
+println("  Dimensions: ", ndims(spec))
+
+# Compile the function using the spec
+compiled_f_sum = Reactant.compile(f_sum, (spec,))
+println("✓ Function compiled successfully")
+
+# Now execute with actual data
+x_actual = Reactant.ConcreteRArray(rand(Float32, 10, 20))
+result = compiled_f_sum(x_actual)
+println("Result: ", result, " (type: ", typeof(result), ")")
+println()
+
+# Example 2: Multiple arguments
+println("Example 2: Multiple arguments with ShapeDtypeStruct")
+println("=" ^ 60)
+
+f_add(x, y) = x .+ y
+
+spec1 = Reactant.ShapeDtypeStruct((5, 5), Float64)
+spec2 = Reactant.ShapeDtypeStruct((5, 5), Float64)
+
+compiled_f_add = Reactant.compile(f_add, (spec1, spec2))
+println("✓ Function with 2 arguments compiled")
+
+x_data = Reactant.ConcreteRArray(rand(Float64, 5, 5))
+y_data = Reactant.ConcreteRArray(rand(Float64, 5, 5))
+result_add = compiled_f_add(x_data, y_data)
+println("Result shape: ", size(result_add))
+println()
+
+# Example 3: Different dtypes
+println("Example 3: Compilation with different dtypes")
+println("=" ^ 60)
+
+f_sin(x) = sin.(x)
+
+for dtype in [Float32, Float64]
+    spec = Reactant.ShapeDtypeStruct((100,), dtype)
+    compiled = Reactant.compile(f_sin, (spec,))
+    
+    x = Reactant.ConcreteRArray(rand(dtype, 100))
+    result = compiled(x)
+    println("✓ Compiled and ran for dtype: ", dtype)
+end
+println()
+
+# Example 4: Benefits demonstration
+println("Example 4: Memory efficiency")
+println("=" ^ 60)
+
+# For very large arrays, you can compile without allocating the full array:
+large_spec = Reactant.ShapeDtypeStruct((10000, 10000), Float32)
+println("Created spec for large array: ", size(large_spec))
+println("  This doesn't allocate ", prod(size(large_spec)) * sizeof(Float32) / 1e9, " GB of memory!")
+
+# Compile a function for this large array
+f_large(x) = sum(x .* x)
+compiled_large = Reactant.compile(f_large, (large_spec,))
+println("✓ Compiled function for large array without allocating memory")
+println()
+
+println("All examples completed successfully!")
+println()
+println("Key Takeaways:")
+println("1. ShapeDtypeStruct allows compilation without data allocation")
+println("2. Same compiled function can be used with actual ConcreteRArray data")
+println("3. Useful for large arrays and rapid prototyping")
+println("4. Similar API to JAX's ShapeDtypeStruct")

From 1652b8f2457ea09d6dc53be34460c194183d2b0b Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 12 Feb 2026 15:42:09 +0000
Subject: [PATCH 5/8] Address code review: use proper exception types and
 rename example file

Co-authored-by: avik-pal <30564094+avik-pal@users.noreply.github.com>
---
 ...xample.jl => shape_dtype_struct_example.jl} |  0
 src/Tracing.jl                                 | 18 +++++++++++++++---
 2 files changed, 15 insertions(+), 3 deletions(-)
 rename examples/{shapedtype_example.jl => shape_dtype_struct_example.jl} (100%)

diff --git a/examples/shapedtype_example.jl b/examples/shape_dtype_struct_example.jl
similarity index 100%
rename from examples/shapedtype_example.jl
rename to examples/shape_dtype_struct_example.jl
diff --git a/src/Tracing.jl b/src/Tracing.jl
index 2536ceae46..5b26d6cfb1 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -1415,15 +1415,27 @@ Base.@nospecializeinfer function make_tracer(
     kwargs...,
 ) where {T,N}
     if mode == TracedToTypes
-        throw("Cannot have ShapeDtypeStruct as function call argument.")
+        throw(
+            ArgumentError(
+                "ShapeDtypeStruct cannot be used as a function call argument; it is only valid for compilation signatures."
+            ),
+        )
     end
     if mode == ArrayToConcrete
-        throw("Cannot convert ShapeDtypeStruct to ConcreteRArray. ShapeDtypeStruct is only for compilation signatures.")
+        throw(
+            ErrorException(
+                "Cannot convert ShapeDtypeStruct to ConcreteRArray. ShapeDtypeStruct is only for compilation signatures."
+            ),
+        )
     end
     # ShapeDtypeStruct behaves like ConcreteToTraced mode - creates a TracedRArray without data
     # Accept both ConcreteToTraced and TracedSetPath modes
     if mode != ConcreteToTraced && mode != TracedSetPath
-        throw("ShapeDtypeStruct can only be used with ConcreteToTraced or TracedSetPath mode, got $mode")
+        throw(
+            ArgumentError(
+                "ShapeDtypeStruct can only be used with ConcreteToTraced or TracedSetPath mode, got $mode"
+            ),
+        )
     end
     haskey(seen, prev) && return seen[prev]::TracedRArray{T,N}
     res = TracedRArray{T,N}((path,), nothing, size(prev))

From 607b13aaffb8562ee3a232e3c84d99eb91b7300e Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Wed, 4 Mar 2026 17:19:00 -0600
Subject: [PATCH 6/8] fix: name

---
 docs/src/api/api.md                    |  2 +-
 examples/shape_dtype_struct_example.jl | 28 +++++++++++------------
 src/Tracing.jl                         | 10 ++++-----
 src/Types.jl                           | 31 +++++++++++++-------------
 test/core/compile.jl                   | 26 ++++++++++-----------
 5 files changed, 49 insertions(+), 48 deletions(-)

diff --git a/docs/src/api/api.md b/docs/src/api/api.md
index edef9a44fa..728b47806b 100644
--- a/docs/src/api/api.md
+++ b/docs/src/api/api.md
@@ -32,7 +32,7 @@ Reactant.to_rarray
 ```@docs
 ConcreteRArray
 ConcreteRNumber
-ShapeDtypeStruct
+ShapedRArray
 ```
 
 ## Inspect Generated HLO
diff --git a/examples/shape_dtype_struct_example.jl b/examples/shape_dtype_struct_example.jl
index 24d4a11b89..ff764901de 100644
--- a/examples/shape_dtype_struct_example.jl
+++ b/examples/shape_dtype_struct_example.jl
@@ -1,16 +1,16 @@
-# Example: Using ShapeDtypeStruct for compilation without concrete arrays
+# Example: Using ShapedRArray for compilation without concrete arrays
 #
-# This example demonstrates how to use Reactant.ShapeDtypeStruct to compile
+# This example demonstrates how to use Reactant.ShapedRArray to compile
 # functions without having to construct full ConcreteRArray instances with
 # actual data. This is useful for:
 # 1. Faster compilation when you only need shape/dtype information
 # 2. Memory efficiency when working with large arrays
-# 3. Similar workflow to JAX's ShapeDtypeStruct
+# 3. Similar workflow to JAX's ShapedRArray
 
 using Reactant
 
 # Example 1: Basic usage with a simple function
-println("Example 1: Basic compilation with ShapeDtypeStruct")
+println("Example 1: Basic compilation with ShapedRArray")
 println("=" ^ 60)
 
 # Define a simple function that sums an array
@@ -19,9 +19,9 @@ f_sum(x) = sum(x)
 # Instead of creating a full ConcreteRArray:
 # x = Reactant.ConcreteRArray(rand(Float32, 10, 20))  # This allocates memory!
 
-# Use ShapeDtypeStruct to specify only shape and dtype:
-spec = Reactant.ShapeDtypeStruct((10, 20), Float32)
-println("Created ShapeDtypeStruct: ", spec)
+# Use ShapedRArray to specify only shape and dtype:
+spec = Reactant.ShapedRArray((10, 20), Float32)
+println("Created ShapedRArray: ", spec)
 println("  Shape: ", size(spec))
 println("  Element type: ", eltype(spec))
 println("  Dimensions: ", ndims(spec))
@@ -37,13 +37,13 @@ println("Result: ", result, " (type: ", typeof(result), ")")
 println()
 
 # Example 2: Multiple arguments
-println("Example 2: Multiple arguments with ShapeDtypeStruct")
+println("Example 2: Multiple arguments with ShapedRArray")
 println("=" ^ 60)
 
 f_add(x, y) = x .+ y
 
-spec1 = Reactant.ShapeDtypeStruct((5, 5), Float64)
-spec2 = Reactant.ShapeDtypeStruct((5, 5), Float64)
+spec1 = Reactant.ShapedRArray((5, 5), Float64)
+spec2 = Reactant.ShapedRArray((5, 5), Float64)
 
 compiled_f_add = Reactant.compile(f_add, (spec1, spec2))
 println("✓ Function with 2 arguments compiled")
@@ -61,7 +61,7 @@ println("=" ^ 60)
 f_sin(x) = sin.(x)
 
 for dtype in [Float32, Float64]
-    spec = Reactant.ShapeDtypeStruct((100,), dtype)
+    spec = Reactant.ShapedRArray((100,), dtype)
     compiled = Reactant.compile(f_sin, (spec,))
     
     x = Reactant.ConcreteRArray(rand(dtype, 100))
@@ -75,7 +75,7 @@ println("Example 4: Memory efficiency")
 println("=" ^ 60)
 
 # For very large arrays, you can compile without allocating the full array:
-large_spec = Reactant.ShapeDtypeStruct((10000, 10000), Float32)
+large_spec = Reactant.ShapedRArray((10000, 10000), Float32)
 println("Created spec for large array: ", size(large_spec))
 println("  This doesn't allocate ", prod(size(large_spec)) * sizeof(Float32) / 1e9, " GB of memory!")
 
@@ -88,7 +88,7 @@ println()
 println("All examples completed successfully!")
 println()
 println("Key Takeaways:")
-println("1. ShapeDtypeStruct allows compilation without data allocation")
+println("1. ShapedRArray allows compilation without data allocation")
 println("2. Same compiled function can be used with actual ConcreteRArray data")
 println("3. Useful for large arrays and rapid prototyping")
-println("4. Similar API to JAX's ShapeDtypeStruct")
+println("4. Similar API to JAX's ShapedRArray")
diff --git a/src/Tracing.jl b/src/Tracing.jl
index 5b26d6cfb1..cbfbe2cd78 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -1409,7 +1409,7 @@ end
 
 Base.@nospecializeinfer function make_tracer(
     seen,
-    @nospecialize(prev::ShapeDtypeStruct{T,N}),
+    @nospecialize(prev::ShapedRArray{T,N}),
     @nospecialize(path),
     mode;
     kwargs...,
@@ -1417,23 +1417,23 @@ Base.@nospecializeinfer function make_tracer(
     if mode == TracedToTypes
         throw(
             ArgumentError(
-                "ShapeDtypeStruct cannot be used as a function call argument; it is only valid for compilation signatures."
+                "ShapedRArray cannot be used as a function call argument; it is only valid for compilation signatures."
             ),
         )
     end
     if mode == ArrayToConcrete
         throw(
             ErrorException(
-                "Cannot convert ShapeDtypeStruct to ConcreteRArray. ShapeDtypeStruct is only for compilation signatures."
+                "Cannot convert ShapedRArray to ConcreteRArray. ShapedRArray is only for compilation signatures."
             ),
         )
     end
-    # ShapeDtypeStruct behaves like ConcreteToTraced mode - creates a TracedRArray without data
+    # ShapedRArray behaves like ConcreteToTraced mode - creates a TracedRArray without data
     # Accept both ConcreteToTraced and TracedSetPath modes
     if mode != ConcreteToTraced && mode != TracedSetPath
         throw(
             ArgumentError(
-                "ShapeDtypeStruct can only be used with ConcreteToTraced or TracedSetPath mode, got $mode"
+                "ShapedRArray can only be used with ConcreteToTraced or TracedSetPath mode, got $mode"
             ),
         )
     end
diff --git a/src/Types.jl b/src/Types.jl
index a2a3e0ca93..ecc549acf9 100644
--- a/src/Types.jl
+++ b/src/Types.jl
@@ -135,13 +135,13 @@ const AnyTracedRVector{T} = AnyTracedRArray{T,1}
 const AnyTracedRMatrix{T} = AnyTracedRArray{T,2}
 const AnyTracedRVecOrMat{T} = Union{AnyTracedRVector{T},AnyTracedRMatrix{T}}
 
-## ShapeDtypeStruct
+## ShapedRArray
 """
-    ShapeDtypeStruct{T,N}(shape::NTuple{N,Int})
-    ShapeDtypeStruct(shape::NTuple{N,Int}, dtype::Type{T}) where {T,N}
+    ShapedRArray{T,N}(shape::NTuple{N,Int})
+    ShapedRArray(shape::NTuple{N,Int}, dtype::Type{T}) where {T,N}
 
 Lightweight structure that specifies the shape and element type (dtype) of an array
-without allocating the actual array data. Similar to JAX's `ShapeDtypeStruct`.
+without allocating the actual array data. Similar to JAX's `ShapedRArray`.
 
 This is useful for compiling functions without constructing the full `ConcreteRArray`,
 which can save memory and improve compilation performance.
@@ -149,7 +149,7 @@ which can save memory and improve compilation performance.
 # Examples
 ```julia
 # Specify shape and dtype for a 2D array
-spec = Reactant.ShapeDtypeStruct((10, 20), Float32)
+spec = Reactant.ShapedRArray((10, 20), Float32)
 
 # Compile a function using just the spec
 f(x) = sum(x)
@@ -162,27 +162,28 @@ result = compiled_f(x)
 
 See also: [`compile`](@ref), [`ConcreteRArray`](@ref)
 """
-struct ShapeDtypeStruct{T,N}
+struct ShapedRArray{T,N}
     shape::NTuple{N,Int}
+    # TODO: Sharding
 
-    function ShapeDtypeStruct{T,N}(shape::NTuple{N,Int}) where {T,N}
+    function ShapedRArray{T,N}(shape::NTuple{N,Int}) where {T,N}
         return new{T,N}(shape)
     end
 end
 
-function ShapeDtypeStruct(shape::NTuple{N,Int}, dtype::Type{T}) where {T,N}
-    return ShapeDtypeStruct{T,N}(shape)
+function ShapedRArray(shape::NTuple{N,Int}, dtype::Type{T}) where {T,N}
+    return ShapedRArray{T,N}(shape)
 end
 
-function ShapeDtypeStruct(shape::Tuple{Vararg{Integer}}, dtype::Type{T}) where {T}
-    return ShapeDtypeStruct(map(Int, shape), dtype)
+function ShapedRArray(shape::Tuple{Vararg{Integer}}, dtype::Type{T}) where {T}
+    return ShapedRArray(map(Int, shape), dtype)
 end
 
-Base.size(x::ShapeDtypeStruct) = x.shape
-Base.ndims(::ShapeDtypeStruct{T,N}) where {T,N} = N
-Base.eltype(::ShapeDtypeStruct{T}) where {T} = T
+Base.size(x::ShapedRArray) = x.shape
+Base.ndims(::ShapedRArray{T,N}) where {T,N} = N
+Base.eltype(::ShapedRArray{T}) where {T} = T
 
-@leaf ShapeDtypeStruct
+@leaf ShapedRArray
 
 # Concrete Types
 ## ConcretePJRTNumber
diff --git a/test/core/compile.jl b/test/core/compile.jl
index 1993ff96b8..55fe754a8f 100644
--- a/test/core/compile.jl
+++ b/test/core/compile.jl
@@ -641,13 +641,13 @@ end
     @test Array(y[:x]) ≈ [2.0f0, 3.0f0]
 end
 
-@testset "ShapeDtypeStruct compilation" begin
-    @testset "Basic compilation with ShapeDtypeStruct" begin
+@testset "ShapedRArray compilation" begin
+    @testset "Basic compilation with ShapedRArray" begin
         # Define a simple function
         f(x) = sum(x)
         
-        # Compile using ShapeDtypeStruct instead of ConcreteRArray
-        spec = Reactant.ShapeDtypeStruct((10, 20), Float32)
+        # Compile using ShapedRArray instead of ConcreteRArray
+        spec = Reactant.ShapedRArray((10, 20), Float32)
         compiled_f = Reactant.compile(f, (spec,))
         
         # Execute with actual data
@@ -658,11 +658,11 @@ end
         @test result ≈ sum(Array(x))
     end
     
-    @testset "Multiple arguments with ShapeDtypeStruct" begin
+    @testset "Multiple arguments with ShapedRArray" begin
         f(x, y) = x .+ y
         
-        spec1 = Reactant.ShapeDtypeStruct((5, 5), Float64)
-        spec2 = Reactant.ShapeDtypeStruct((5, 5), Float64)
+        spec1 = Reactant.ShapedRArray((5, 5), Float64)
+        spec2 = Reactant.ShapedRArray((5, 5), Float64)
         compiled_f = Reactant.compile(f, (spec1, spec2))
         
         x = Reactant.ConcreteRArray(rand(Float64, 5, 5))
@@ -673,11 +673,11 @@ end
         @test result ≈ Array(x) .+ Array(y)
     end
     
-    @testset "ShapeDtypeStruct with different dtypes" begin
+    @testset "ShapedRArray with different dtypes" begin
         f(x) = sin.(x)
         
         for dtype in [Float32, Float64]
-            spec = Reactant.ShapeDtypeStruct((10,), dtype)
+            spec = Reactant.ShapedRArray((10,), dtype)
             compiled_f = Reactant.compile(f, (spec,))
             
             x = Reactant.ConcreteRArray(rand(dtype, 10))
@@ -688,19 +688,19 @@ end
         end
     end
     
-    @testset "ShapeDtypeStruct constructor variations" begin
+    @testset "ShapedRArray constructor variations" begin
         # Test different constructor forms
-        spec1 = Reactant.ShapeDtypeStruct{Float32,2}((3, 4))
+        spec1 = Reactant.ShapedRArray{Float32,2}((3, 4))
         @test size(spec1) == (3, 4)
         @test eltype(spec1) == Float32
         @test ndims(spec1) == 2
         
-        spec2 = Reactant.ShapeDtypeStruct((3, 4), Float32)
+        spec2 = Reactant.ShapedRArray((3, 4), Float32)
         @test size(spec2) == (3, 4)
         @test eltype(spec2) == Float32
         
         # Test with integer tuple (not Int tuple)
-        spec3 = Reactant.ShapeDtypeStruct((3, 4), Float64)
+        spec3 = Reactant.ShapedRArray((3, 4), Float64)
         @test size(spec3) == (3, 4)
         @test eltype(spec3) == Float64
     end

From bd5a97b711c070900e403a7168c7807e142594f4 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Wed, 4 Mar 2026 18:54:23 -0600
Subject: [PATCH 7/8] fix: general fixes to compilation pipeline

---
 docs/src/api/api.md                    |  3 +-
 examples/shape_dtype_struct_example.jl | 94 --------------------------
 src/Compiler.jl                        |  4 +-
 src/ConcreteRArray.jl                  | 12 ++++
 src/Reactant.jl                        |  8 +--
 src/Tracing.jl                         | 76 +++++++++++++++------
 src/Types.jl                           | 61 +++++++++++------
 test/core/compile.jl                   | 61 ++++++++---------
 8 files changed, 146 insertions(+), 173 deletions(-)
 delete mode 100644 examples/shape_dtype_struct_example.jl

diff --git a/docs/src/api/api.md b/docs/src/api/api.md
index 728b47806b..1b9fa1e1df 100644
--- a/docs/src/api/api.md
+++ b/docs/src/api/api.md
@@ -32,7 +32,8 @@ Reactant.to_rarray
 ```@docs
 ConcreteRArray
 ConcreteRNumber
-ShapedRArray
+RArraySpec
+RNumberSpec
 ```
 
 ## Inspect Generated HLO
diff --git a/examples/shape_dtype_struct_example.jl b/examples/shape_dtype_struct_example.jl
deleted file mode 100644
index ff764901de..0000000000
--- a/examples/shape_dtype_struct_example.jl
+++ /dev/null
@@ -1,94 +0,0 @@
-# Example: Using ShapedRArray for compilation without concrete arrays
-#
-# This example demonstrates how to use Reactant.ShapedRArray to compile
-# functions without having to construct full ConcreteRArray instances with
-# actual data. This is useful for:
-# 1. Faster compilation when you only need shape/dtype information
-# 2. Memory efficiency when working with large arrays
-# 3. Similar workflow to JAX's ShapedRArray
-
-using Reactant
-
-# Example 1: Basic usage with a simple function
-println("Example 1: Basic compilation with ShapedRArray")
-println("=" ^ 60)
-
-# Define a simple function that sums an array
-f_sum(x) = sum(x)
-
-# Instead of creating a full ConcreteRArray:
-# x = Reactant.ConcreteRArray(rand(Float32, 10, 20))  # This allocates memory!
-
-# Use ShapedRArray to specify only shape and dtype:
-spec = Reactant.ShapedRArray((10, 20), Float32)
-println("Created ShapedRArray: ", spec)
-println("  Shape: ", size(spec))
-println("  Element type: ", eltype(spec))
-println("  Dimensions: ", ndims(spec))
-
-# Compile the function using the spec
-compiled_f_sum = Reactant.compile(f_sum, (spec,))
-println("✓ Function compiled successfully")
-
-# Now execute with actual data
-x_actual = Reactant.ConcreteRArray(rand(Float32, 10, 20))
-result = compiled_f_sum(x_actual)
-println("Result: ", result, " (type: ", typeof(result), ")")
-println()
-
-# Example 2: Multiple arguments
-println("Example 2: Multiple arguments with ShapedRArray")
-println("=" ^ 60)
-
-f_add(x, y) = x .+ y
-
-spec1 = Reactant.ShapedRArray((5, 5), Float64)
-spec2 = Reactant.ShapedRArray((5, 5), Float64)
-
-compiled_f_add = Reactant.compile(f_add, (spec1, spec2))
-println("✓ Function with 2 arguments compiled")
-
-x_data = Reactant.ConcreteRArray(rand(Float64, 5, 5))
-y_data = Reactant.ConcreteRArray(rand(Float64, 5, 5))
-result_add = compiled_f_add(x_data, y_data)
-println("Result shape: ", size(result_add))
-println()
-
-# Example 3: Different dtypes
-println("Example 3: Compilation with different dtypes")
-println("=" ^ 60)
-
-f_sin(x) = sin.(x)
-
-for dtype in [Float32, Float64]
-    spec = Reactant.ShapedRArray((100,), dtype)
-    compiled = Reactant.compile(f_sin, (spec,))
-    
-    x = Reactant.ConcreteRArray(rand(dtype, 100))
-    result = compiled(x)
-    println("✓ Compiled and ran for dtype: ", dtype)
-end
-println()
-
-# Example 4: Benefits demonstration
-println("Example 4: Memory efficiency")
-println("=" ^ 60)
-
-# For very large arrays, you can compile without allocating the full array:
-large_spec = Reactant.ShapedRArray((10000, 10000), Float32)
-println("Created spec for large array: ", size(large_spec))
-println("  This doesn't allocate ", prod(size(large_spec)) * sizeof(Float32) / 1e9, " GB of memory!")
-
-# Compile a function for this large array
-f_large(x) = sum(x .* x)
-compiled_large = Reactant.compile(f_large, (large_spec,))
-println("✓ Compiled function for large array without allocating memory")
-println()
-
-println("All examples completed successfully!")
-println()
-println("Key Takeaways:")
-println("1. ShapedRArray allows compilation without data allocation")
-println("2. Same compiled function can be used with actual ConcreteRArray data")
-println("3. Useful for large arrays and rapid prototyping")
-println("4. Similar API to JAX's ShapedRArray")
diff --git a/src/Compiler.jl b/src/Compiler.jl
index 4b62f914b8..d207796c25 100644
--- a/src/Compiler.jl
+++ b/src/Compiler.jl
@@ -3947,7 +3947,9 @@ 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
+                continue
+            end
             buffer = k.data isa Tuple ? only(k.data) : k.data
             push!(devices_list, XLA.device(buffer))
         end
diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 6a57d31ba8..8015269c7f 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -286,6 +286,18 @@ 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.print_array(io::IO, X::Union{AnyConcretePJRTArray,AnyConcreteIFRTArray})
     if isempty(X)
         print(io, "<Empty Buffer eltype $(eltype(X)) of size $(size(X))>")
diff --git a/src/Reactant.jl b/src/Reactant.jl
index 219eb60288..ab0c36e543 100644
--- a/src/Reactant.jl
+++ b/src/Reactant.jl
@@ -124,9 +124,8 @@ 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)
 
-# For ShapeDtypeStruct
-unwrapped_eltype(::Type{ShapeDtypeStruct{T,N}}) where {T,N} = T
-unwrapped_eltype(::ShapeDtypeStruct{T,N}) where {T,N} = 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)
@@ -293,7 +292,8 @@ export ConcreteRArray,
     ConcretePJRTNumber,
     ConcreteIFRTArray,
     ConcreteIFRTNumber,
-    ShapeDtypeStruct,
+    RArraySpec,
+    RNumberSpec,
     @compile,
     @code_hlo,
     @code_mhlo,
diff --git a/src/Tracing.jl b/src/Tracing.jl
index cbfbe2cd78..926c2edcc8 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -3,6 +3,7 @@
     TracedTrack = 2
     TracedToConcrete = 3
     ArrayToConcrete = 4
+    # TODO: Array to Specification
     TracedSetPath = 5
     TracedToTypes = 6
     NoStopTracedTrack = 7
@@ -329,6 +330,45 @@ Base.@nospecializeinfer function traced_type_inner(
     end
 end
 
+Base.@nospecializeinfer function traced_type_inner(
+    @nospecialize(T::Type{<:RArraySpec}),
+    seen,
+    @nospecialize(mode::TraceMode),
+    @nospecialize(track_numbers::Type),
+    @nospecialize(ndevices),
+    @nospecialize(runtime)
+)
+    if mode == ConcreteToTraced
+        T´ = Base.unwrap_unionall(T)
+        T, N = T´.parameters
+        T´´ = TracedRArray{T,N}
+        T_ret = N isa Core.TypeVar ? UnionAll(N, T´´) : T´´
+        T_ret2 = T isa Core.TypeVar ? UnionAll(T, T_ret) : T_ret
+        return T_ret2
+    else
+        throw("Unsupported mode: $mode")
+    end
+end
+
+Base.@nospecializeinfer function traced_type_inner(
+    @nospecialize(T::Type{<:RNumberSpec}),
+    seen,
+    @nospecialize(mode::TraceMode),
+    @nospecialize(track_numbers::Type),
+    @nospecialize(ndevices),
+    @nospecialize(runtime)
+)
+    if mode == ConcreteToTraced
+        T´ = Base.unwrap_unionall(T)
+        T = T´.parameters
+        T´´ = TracedRNumber{T}
+        T_ret = T isa Core.TypeVar ? UnionAll(T, T´´) : T´´
+        return T_ret
+    else
+        throw("Unsupported mode: $mode")
+    end
+end
+
 Base.@nospecializeinfer function traced_type_inner(
     @nospecialize(T::Type{MissingTracedValue}),
     seen,
@@ -1408,37 +1448,33 @@ Base.@nospecializeinfer function make_tracer(
 end
 
 Base.@nospecializeinfer function make_tracer(
-    seen,
-    @nospecialize(prev::ShapedRArray{T,N}),
-    @nospecialize(path),
-    mode;
-    kwargs...,
+    seen, @nospecialize(prev::RArraySpec{T,N}), @nospecialize(path), mode; kwargs...
 ) where {T,N}
-    if mode == TracedToTypes
+    if mode != ConcreteToTraced
         throw(
             ArgumentError(
-                "ShapedRArray cannot be used as a function call argument; it is only valid for compilation signatures."
-            ),
-        )
-    end
-    if mode == ArrayToConcrete
-        throw(
-            ErrorException(
-                "Cannot convert ShapedRArray to ConcreteRArray. ShapedRArray is only for compilation signatures."
+                "RArraySpec can only be used with ConcreteToTraced mode, got $mode"
             ),
         )
     end
-    # ShapedRArray behaves like ConcreteToTraced mode - creates a TracedRArray without data
-    # Accept both ConcreteToTraced and TracedSetPath modes
-    if mode != ConcreteToTraced && mode != TracedSetPath
+    haskey(seen, prev) && return seen[prev]::TracedRArray{T,N}
+    res = TracedRArray{T,N}((path,), nothing, size(prev))
+    seen[prev] = res
+    return res
+end
+
+Base.@nospecializeinfer function make_tracer(
+    seen, @nospecialize(prev::RNumberSpec{T}), @nospecialize(path), mode; kwargs...
+) where {T}
+    if mode != ConcreteToTraced
         throw(
             ArgumentError(
-                "ShapedRArray can only be used with ConcreteToTraced or TracedSetPath mode, got $mode"
+                "RNumberSpec can only be used with ConcreteToTraced mode, got $mode"
             ),
         )
     end
-    haskey(seen, prev) && return seen[prev]::TracedRArray{T,N}
-    res = TracedRArray{T,N}((path,), nothing, size(prev))
+    haskey(seen, prev) && return seen[prev]::TracedRNumber{T}
+    res = TracedRNumber{T}((path,), nothing)
     seen[prev] = res
     return res
 end
diff --git a/src/Types.jl b/src/Types.jl
index ecc549acf9..8f860eb7ed 100644
--- a/src/Types.jl
+++ b/src/Types.jl
@@ -135,13 +135,13 @@ const AnyTracedRVector{T} = AnyTracedRArray{T,1}
 const AnyTracedRMatrix{T} = AnyTracedRArray{T,2}
 const AnyTracedRVecOrMat{T} = Union{AnyTracedRVector{T},AnyTracedRMatrix{T}}
 
-## ShapedRArray
+## RArraySpec
 """
-    ShapedRArray{T,N}(shape::NTuple{N,Int})
-    ShapedRArray(shape::NTuple{N,Int}, dtype::Type{T}) where {T,N}
+    RArraySpec{T,N}(shape::NTuple{N,Int})
+    RArraySpec{T}(shape::NTuple{N,Int})
 
 Lightweight structure that specifies the shape and element type (dtype) of an array
-without allocating the actual array data. Similar to JAX's `ShapedRArray`.
+without allocating the actual array data. Similar to JAX's `RArraySpec`.
 
 This is useful for compiling functions without constructing the full `ConcreteRArray`,
 which can save memory and improve compilation performance.
@@ -149,7 +149,7 @@ which can save memory and improve compilation performance.
 # Examples
 ```julia
 # Specify shape and dtype for a 2D array
-spec = Reactant.ShapedRArray((10, 20), Float32)
+spec = Reactant.RArraySpec{Float32}((10, 20))
 
 # Compile a function using just the spec
 f(x) = sum(x)
@@ -162,28 +162,51 @@ result = compiled_f(x)
 
 See also: [`compile`](@ref), [`ConcreteRArray`](@ref)
 """
-struct ShapedRArray{T,N}
+struct RArraySpec{T,N} <: RArray{T,N}
     shape::NTuple{N,Int}
     # TODO: Sharding
-
-    function ShapedRArray{T,N}(shape::NTuple{N,Int}) where {T,N}
-        return new{T,N}(shape)
-    end
 end
 
-function ShapedRArray(shape::NTuple{N,Int}, dtype::Type{T}) where {T,N}
-    return ShapedRArray{T,N}(shape)
-end
+RArraySpec{T}(shape::NTuple{N,Int}) where {T,N} = RArraySpec{T,N}(shape)
+
+Base.size(x::RArraySpec) = x.shape
+Base.ndims(::RArraySpec{T,N}) where {T,N} = N
+Base.eltype(::RArraySpec{T}) where {T} = T
+
+@leaf RArraySpec
+
+"""
+    RNumberSpec{T}()
+
+Lightweight structure that specifies the element type (dtype) of a number
+without allocating the actual number data. Similar to JAX's `RNumberSpec`.
+
+This is useful for compiling functions without constructing the full `ConcreteRNumber`,
+which can save memory and improve compilation performance.
+
+# Examples
+```julia
+# Specify dtype for a number
+spec = Reactant.RNumberSpec{Float32}()
+
+# Compile a function using just the spec
+f(x) = x + 1
+compiled_f = Reactant.compile(f, (spec,))
 
-function ShapedRArray(shape::Tuple{Vararg{Integer}}, dtype::Type{T}) where {T}
-    return ShapedRArray(map(Int, shape), dtype)
+# Execute with actual data
+x = Reactant.ConcreteRNumber(1.0f0)
+result = compiled_f(x)
+```
+
+See also: [`compile`](@ref), [`ConcreteRNumber`](@ref)
+"""
+struct RNumberSpec{T} <: RNumber{T}
+    # TODO: Sharding
 end
 
-Base.size(x::ShapedRArray) = x.shape
-Base.ndims(::ShapedRArray{T,N}) where {T,N} = N
-Base.eltype(::ShapedRArray{T}) where {T} = T
+Base.eltype(::RNumberSpec{T}) where {T} = T
 
-@leaf ShapedRArray
+@leaf RNumberSpec
 
 # Concrete Types
 ## ConcretePJRTNumber
diff --git a/test/core/compile.jl b/test/core/compile.jl
index 55fe754a8f..2e8e6613e3 100644
--- a/test/core/compile.jl
+++ b/test/core/compile.jl
@@ -641,66 +641,59 @@ end
     @test Array(y[:x]) ≈ [2.0f0, 3.0f0]
 end
 
-@testset "ShapedRArray compilation" begin
-    @testset "Basic compilation with ShapedRArray" begin
-        # Define a simple function
-        f(x) = sum(x)
-        
-        # Compile using ShapedRArray instead of ConcreteRArray
-        spec = Reactant.ShapedRArray((10, 20), Float32)
-        compiled_f = Reactant.compile(f, (spec,))
-        
+@testset "RArraySpec compilation" begin
+    @testset "Basic compilation with RArraySpec" begin
+        # Compile using RArraySpec instead of ConcreteRArray
+        spec = RArraySpec{Float32}((10, 20))
+        compiled_f = Reactant.compile(sum, (spec,))
+
         # Execute with actual data
         x = Reactant.ConcreteRArray(rand(Float32, 10, 20))
         result = compiled_f(x)
-        
+
         @test result isa Reactant.ConcreteRNumber{Float32}
         @test result ≈ sum(Array(x))
     end
-    
-    @testset "Multiple arguments with ShapedRArray" begin
-        f(x, y) = x .+ y
-        
-        spec1 = Reactant.ShapedRArray((5, 5), Float64)
-        spec2 = Reactant.ShapedRArray((5, 5), Float64)
-        compiled_f = Reactant.compile(f, (spec1, spec2))
-        
+
+    @testset "Multiple arguments with RArraySpec" begin
+        spec1 = Reactant.RArraySpec((5, 5), Float64)
+        spec2 = Reactant.RArraySpec((5, 5), Float64)
+        compiled_f = Reactant.compile(.+, (spec1, spec2))
+
         x = Reactant.ConcreteRArray(rand(Float64, 5, 5))
         y = Reactant.ConcreteRArray(rand(Float64, 5, 5))
         result = compiled_f(x, y)
-        
+
         @test result isa Reactant.ConcreteRArray{Float64,2}
         @test result ≈ Array(x) .+ Array(y)
     end
-    
-    @testset "ShapedRArray with different dtypes" begin
-        f(x) = sin.(x)
-        
+
+    @testset "RArraySpec with different dtypes" begin
         for dtype in [Float32, Float64]
-            spec = Reactant.ShapedRArray((10,), dtype)
-            compiled_f = Reactant.compile(f, (spec,))
-            
+            spec = Reactant.RArraySpec((10,), dtype)
+            compiled_f = Reactant.compile(Base.BroadcastFunction(sin), (spec,))
+
             x = Reactant.ConcreteRArray(rand(dtype, 10))
             result = compiled_f(x)
-            
+
             @test result isa Reactant.ConcreteRArray{dtype,1}
             @test result ≈ sin.(Array(x))
         end
     end
-    
-    @testset "ShapedRArray constructor variations" begin
+
+    @testset "RArraySpec constructor variations" begin
         # Test different constructor forms
-        spec1 = Reactant.ShapedRArray{Float32,2}((3, 4))
+        spec1 = Reactant.RArraySpec{Float32,2}((3, 4))
         @test size(spec1) == (3, 4)
         @test eltype(spec1) == Float32
         @test ndims(spec1) == 2
-        
-        spec2 = Reactant.ShapedRArray((3, 4), Float32)
+
+        spec2 = Reactant.RArraySpec((3, 4), Float32)
         @test size(spec2) == (3, 4)
         @test eltype(spec2) == Float32
-        
+
         # Test with integer tuple (not Int tuple)
-        spec3 = Reactant.ShapedRArray((3, 4), Float64)
+        spec3 = Reactant.RArraySpec((3, 4), Float64)
         @test size(spec3) == (3, 4)
         @test eltype(spec3) == Float64
     end

From cee199240634261cd4042c5fbd3c528ea3503a43 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Wed, 4 Mar 2026 19:39:41 -0600
Subject: [PATCH 8/8] feat: construct specification from structures

---
 src/Compiler.jl       |  5 ++-
 src/ConcreteRArray.jl |  5 +++
 src/Tracing.jl        | 80 ++++++++++++++++++++++++++++++++++++-------
 src/Types.jl          | 14 +++++++-
 4 files changed, 89 insertions(+), 15 deletions(-)

diff --git a/src/Compiler.jl b/src/Compiler.jl
index d207796c25..9f0b885330 100644
--- a/src/Compiler.jl
+++ b/src/Compiler.jl
@@ -3947,7 +3947,10 @@ 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
-            if !(v isa TracedRArray || v isa TracedRNumber) || k isa Reactant.RArraySpec
+            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
diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 8015269c7f..5bb4b96f7a 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -298,6 +298,11 @@ function Base.print_array(io::IO, ::RArraySpec)
     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))>")
diff --git a/src/Tracing.jl b/src/Tracing.jl
index 926c2edcc8..4fabe2799c 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -3,11 +3,11 @@
     TracedTrack = 2
     TracedToConcrete = 3
     ArrayToConcrete = 4
-    # TODO: Array to Specification
     TracedSetPath = 5
     TracedToTypes = 6
     NoStopTracedTrack = 7
     TracedToJAX = 8
+    ArrayToSpec = 9
 end
 
 function convert_to_jax_dtype_struct end
@@ -70,6 +70,8 @@ Base.@nospecializeinfer function traced_type_inner(
         else
             error("Unsupported runtime $runtime")
         end
+    elseif mode == ArrayToSpec && T <: track_numbers
+        return RNumberSpec{T}
     elseif (mode == NoStopTracedTrack || mode == TracedTrack || mode == TracedSetPath) &&
         T <: track_numbers
         return TracedRNumber{T}
@@ -239,7 +241,7 @@ Base.@nospecializeinfer function traced_type_inner(
     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)}
         end
@@ -522,6 +524,8 @@ Base.@nospecializeinfer function traced_type_inner(
             else
                 error("Unsupported runtime $runtime")
             end
+        elseif mode == ArrayToSpec && T <: ReactantPrimitive
+            A_wrapper = RArraySpec
         end
 
         # WARN replacing typevars first is required to construct the UnionAlls correctly
@@ -544,6 +548,8 @@ Base.@nospecializeinfer function traced_type_inner(
                 end
             end
             error("Unsupported runtime $runtime")
+        elseif mode == ArrayToSpec && T <: ReactantPrimitive
+            return RArraySpec{T,N}
         else
             return Array{
                 traced_type_inner(T, seen, mode, track_numbers, ndevices, runtime),N
@@ -560,7 +566,7 @@ Base.@nospecializeinfer function traced_type_inner(
     @nospecialize(ndevices),
     @nospecialize(runtime)
 )
-    if mode == ArrayToConcrete
+    if mode == ArrayToConcrete || mode == ArrayToSpec
         A´ = A isa UnionAll ? Array{Bool} : Array{Bool,ndims(A)}
         return traced_type_inner(A´, seen, mode, track_numbers, ndevices, runtime)
     else
@@ -657,7 +663,7 @@ Base.@nospecializeinfer function traced_type_inner(
     @nospecialize(ndevices),
     @nospecialize(runtime)
 )
-    if mode == ArrayToConcrete
+    if mode == ArrayToConcrete || mode == ArrayToSpec
         return ReactantRNG{
             traced_type_inner(Array{UInt64,1}, seen, mode, track_numbers, ndevices, runtime)
         }
@@ -1760,6 +1766,8 @@ Base.@nospecializeinfer function make_tracer(
             runtime isa Val{:IFRT} &&
                 return ConcreteIFRTNumber(prev; sharding, device, client)
             error("Unsupported runtime $runtime")
+        elseif mode == ArrayToSpec
+            return RNumberSpec{RT}(; sharding)
         else
             if mode == TracedTrack || mode == NoStopTracedTrack
                 res = TracedRNumber{RT}((path,), broadcast_to_size(prev, ()).mlir_data)
@@ -1852,6 +1860,8 @@ Base.@nospecializeinfer function make_tracer(
             runtime isa Val{:IFRT} &&
                 (return seen[prev] = ConcreteIFRTArray(prev; sharding, device, client))
             error("Unsupported runtime $runtime")
+        elseif mode == ArrayToSpec
+            return seen[prev] = RArraySpec{eltype(RT),ndims(RT)}(size(prev); sharding)
         elseif mode == TracedToTypes
             # Original array can get mutated so we store a copy:
             push!(path, copy(prev))
@@ -1914,7 +1924,7 @@ end
 Base.@nospecializeinfer function make_tracer(
     seen, @nospecialize(prev::BitArray), @nospecialize(path), mode; kwargs...
 )
-    if mode == ArrayToConcrete
+    if mode == ArrayToConcrete || mode == ArrayToSpec
         return make_tracer(seen, Array(prev), path, mode; kwargs...)
     else
         return prev
@@ -1949,6 +1959,8 @@ Base.@nospecializeinfer function make_tracer(
             runtime isa Val{:IFRT} &&
                 (return seen[prev] = ConcreteIFRTArray(prev; sharding, device, client))
             error("Unsupported runtime $runtime")
+        elseif mode == ArrayToSpec
+            return seen[prev] = RArraySpec{eltype(prev),ndims(prev)}(size(prev); sharding)
         elseif mode == TracedToTypes
             # Original array can get mutated so we store a copy:
             push!(path, copy(prev))
@@ -2156,7 +2168,7 @@ end
 Base.@nospecializeinfer function make_tracer(
     seen, @nospecialize(prev::Random.AbstractRNG), @nospecialize(path), mode; kwargs...
 )
-    if mode == ArrayToConcrete
+    if mode == ArrayToConcrete || mode == ArrayToSpec
         TracedRandom.should_warn_if_not_natively_supported(prev)
         return ReactantRNG(
             make_tracer(seen, TracedRandom.make_seed(prev), (path..., 1), mode; kwargs...),
@@ -2167,7 +2179,14 @@ Base.@nospecializeinfer function make_tracer(
 end
 
 """
-    to_rarray(x; track_numbers=false, sharding=NoSharding(), device=nothing, client=nothing, runtime=nothing)
+    to_rarray(
+        x;
+        track_numbers=false,
+        sharding=NoSharding(),
+        device=nothing,
+        client=nothing,
+        runtime=nothing,
+    )
 
 Convert a Julia value `x` into its Reactant equivalent by tracing through the structure.
 Arrays are converted to `ConcreteRArray`, and (optionally) scalar numbers are converted
@@ -2186,6 +2205,9 @@ to `ConcreteRNumber`.
 - `device`: Target device for the resulting array.
 - `client`: XLA client to use.
 - `runtime`: Backend runtime to use (`Val(:PJRT)` or `Val(:IFRT)`).
+- `convert_to_specification::Union{Val{true},Val{false}}=Val{false}()`: Whether to convert
+  the arrays to a specification (RArraySpec/RNumberSpec) instead of a concrete array.
+  This enables AoT compiling reactant functions without allocating the actual arrays.
 
 ## Examples
 
@@ -2214,13 +2236,17 @@ become compile-time constants.
     sharding=Sharding.Sharding.NoSharding(),
     device=nothing,
     client=nothing,
+    convert_to_specification::Union{Val{true},Val{false}}=Val{false}(),
 )
     runtime === nothing && (runtime = XLA.runtime())
     track_numbers isa Bool && (track_numbers = track_numbers ? Number : Union{})
-    return to_rarray_internal(x, track_numbers, sharding, runtime, device, client)
+    return to_rarray_internal(
+        convert_to_specification, x, track_numbers, sharding, runtime, device, client
+    )
 end
 
 @inline function to_rarray_internal(
+    @nospecialize(convert_to_specification),
     @nospecialize(x),
     @nospecialize(track_numbers::Type),
     @nospecialize(sharding),
@@ -2232,7 +2258,7 @@ end
         OrderedIdDict(),
         x,
         (),
-        ArrayToConcrete;
+        convert_to_specification isa Val{true} ? ArrayToSpec : ArrayToConcrete;
         track_numbers,
         sharding,
         runtime,
@@ -2243,6 +2269,7 @@ end
 
 # fast paths avoiding make_tracer
 function to_rarray_internal(
+    @nospecialize(convert_to_specification),
     @nospecialize(::TracedRArray),
     @nospecialize(track_numbers::Type),
     @nospecialize(sharding),
@@ -2254,6 +2281,7 @@ function to_rarray_internal(
 end
 
 @inline function to_rarray_internal(
+    @nospecialize(convert_to_specification),
     @nospecialize(x::ConcretePJRTArray),
     @nospecialize(track_numbers::Type),
     @nospecialize(sharding),
@@ -2265,6 +2293,7 @@ end
 end
 
 @inline function to_rarray_internal(
+    @nospecialize(convert_to_specification),
     @nospecialize(x::ConcreteIFRTArray),
     @nospecialize(track_numbers::Type),
     @nospecialize(sharding),
@@ -2276,6 +2305,7 @@ end
 end
 
 @inline function to_rarray_internal(
+    @nospecialize(convert_to_specification),
     @nospecialize(x::Array{<:ReactantPrimitive}),
     @nospecialize(track_numbers::Type),
     @nospecialize(sharding),
@@ -2283,12 +2313,15 @@ end
     @nospecialize(device),
     @nospecialize(client)
 )
+    convert_to_specification isa Val{true} &&
+        return RArraySpec{eltype(x),ndims(x)}(size(x); sharding)
     runtime isa Val{:PJRT} && return ConcretePJRTArray(x; sharding, device, client)
     runtime isa Val{:IFRT} && return ConcreteIFRTArray(x; sharding, device, client)
     return error("Unsupported runtime $runtime")
 end
 
 @inline function to_rarray_internal(
+    @nospecialize(convert_to_specification),
     @nospecialize(x::Array{T}),
     @nospecialize(track_numbers::Type),
     @nospecialize(sharding),
@@ -2297,6 +2330,8 @@ end
     @nospecialize(client)
 ) where {T<:Number}
     if reactant_primitive(T) !== nothing
+        convert_to_specification isa Val{true} &&
+            return RArraySpec{to_reactant_primitive(T),ndims(x)}(size(x); sharding)
         if runtime isa Val{:PJRT}
             return ConcretePJRTArray(to_reactant_primitive.(x); sharding, device, client)
         elseif runtime isa Val{:IFRT}
@@ -2305,11 +2340,18 @@ end
         error("Unsupported runtime $runtime")
     end
     return @invoke to_rarray_internal(
-        x::Any, track_numbers::Type, sharding, runtime, device, client
+        convert_to_specification,
+        x::Any,
+        track_numbers::Type,
+        sharding,
+        runtime,
+        device,
+        client,
     )
 end
 
 @inline function to_rarray_internal(
+    @nospecialize(convert_to_specification),
     @nospecialize(x::ConcretePJRTNumber),
     @nospecialize(track_numbers::Type),
     @nospecialize(sharding),
@@ -2321,6 +2363,7 @@ end
 end
 
 @inline function to_rarray_internal(
+    @nospecialize(convert_to_specification),
     @nospecialize(x::ConcreteIFRTNumber),
     @nospecialize(track_numbers::Type),
     @nospecialize(sharding),
@@ -2332,6 +2375,7 @@ end
 end
 
 @inline function to_rarray_internal(
+    @nospecialize(convert_to_specification),
     @nospecialize(x::ReactantPrimitive),
     @nospecialize(track_numbers::Type),
     @nospecialize(sharding),
@@ -2340,6 +2384,7 @@ end
     @nospecialize(client)
 )
     if typeof(x) <: track_numbers
+        convert_to_specification isa Val{true} && return RNumberSpec{eltype(x)}(; sharding)
         runtime isa Val{:PJRT} && return ConcretePJRTNumber(x; sharding, device, client)
         runtime isa Val{:IFRT} && return ConcreteIFRTNumber(x; sharding, device, client)
         error("Unsupported runtime $runtime")
@@ -2348,6 +2393,7 @@ end
 end
 
 @inline function to_rarray_internal(
+    @nospecialize(convert_to_specification),
     @nospecialize(x::Number),
     @nospecialize(track_numbers::Type),
     @nospecialize(sharding),
@@ -2356,14 +2402,22 @@ end
     @nospecialize(client)
 )
     if reactant_primitive(typeof(x)) !== nothing
+        convert_to_specification isa Val{true} &&
+            return RNumberSpec{to_reactant_primitive(eltype(x))}(; sharding)
         runtime isa Val{:PJRT} &&
-            return ConcretePJRTArray(to_reactant_primitive(x); sharding, device, client)
+            return ConcretePJRTNumber(to_reactant_primitive(x); sharding, device, client)
         runtime isa Val{:IFRT} &&
-            return ConcreteIFRTArray(to_reactant_primitive(x); sharding, device, client)
+            return ConcreteIFRTNumber(to_reactant_primitive(x); sharding, device, client)
         error("Unsupported runtime $runtime")
     end
     return @invoke to_rarray_internal(
-        x::Any, track_numbers::Type, sharding, runtime, device, client
+        convert_to_specification,
+        x::Any,
+        track_numbers::Type,
+        sharding,
+        runtime,
+        device,
+        client,
     )
 end
 
diff --git a/src/Types.jl b/src/Types.jl
index 8f860eb7ed..c2e24e1924 100644
--- a/src/Types.jl
+++ b/src/Types.jl
@@ -165,9 +165,17 @@ See also: [`compile`](@ref), [`ConcreteRArray`](@ref)
 struct RArraySpec{T,N} <: RArray{T,N}
     shape::NTuple{N,Int}
     # TODO: Sharding
+
+    function RArraySpec{T,N}(
+        shape::NTuple{N,Int}; sharding=Sharding.NoShardInfo()
+    ) where {T,N}
+        return new{T,N}(shape)
+    end
 end
 
-RArraySpec{T}(shape::NTuple{N,Int}) where {T,N} = RArraySpec{T,N}(shape)
+function RArraySpec{T}(shape::NTuple{N,Int}; sharding=Sharding.NoShardInfo()) where {T,N}
+    return RArraySpec{T,N}(shape; sharding)
+end
 
 Base.size(x::RArraySpec) = x.shape
 Base.ndims(::RArraySpec{T,N}) where {T,N} = N
@@ -202,6 +210,10 @@ See also: [`compile`](@ref), [`ConcreteRNumber`](@ref)
 """
 struct RNumberSpec{T} <: RNumber{T}
     # TODO: Sharding
+
+    function RNumberSpec{T}(; sharding=Sharding.NoShardInfo()) where {T}
+        return new{T}()
+    end
 end
 
 Base.eltype(::RNumberSpec{T}) where {T} = T