From f893df2ed8a43ecc1dc2aef3d4be3d26e0e1d285 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 4 Oct 2024 13:44:17 -0400
Subject: [PATCH] feat: ensure size of value in setproperty is same as original
 shape

---
 src/TracedRArray.jl | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
index bc61862..3019926 100644
--- a/src/TracedRArray.jl
+++ b/src/TracedRArray.jl
@@ -17,6 +17,15 @@ mutable struct TracedRArray{T,N} <: RArray{T,N}
     end
 end
 
+function Base.setproperty!(x::TracedRArray, f::Symbol, v)
+    if f === :shape && !isnothing(v)
+        @assert v isa MLIR.IR.Value
+        mlir_size = size(MLIR.IR.type(v))
+        @assert mlir_size == x.shape "Expected shape $(x.shape), got $(mlir_size)"
+    end
+    return setfield!(x, f, v)
+end
+
 const WrappedTracedRArray{T,N} = WrappedArray{T,N,TracedRArray,TracedRArray{T,N}}
 const AnyTracedRArray{T,N} = Union{TracedRArray{T,N},WrappedTracedRArray{T,N}}
 const AnyTracedRScalar{T} = AnyTracedRArray{T,0}