fix: batching

Author: avik-pal

src/TracedRArray.jl

@@ -19,13 +19,6 @@ end
 
 TracedRArray{T,N}(x::TracedRArray{T,N}) where {T,N} = x
 
-function Base.setproperty!(x::TracedRArray, f::Symbol, v)
-    if f === :mlir_data && !isnothing(v)
-        @assert size(MLIR.IR.type(v)) == size(x)
-    end
-    return setfield!(x, f, v)
-end
-
 mutable struct TracedRNumber{T} <: RNumber{T}
     paths::Tuple
     mlir_data::Union{Nothing,MLIR.IR.Value}
@@ -40,13 +33,6 @@ mutable struct TracedRNumber{T} <: RNumber{T}
     end
 end
 
-function Base.setproperty!(x::TracedRNumber, f::Symbol, v)
-    if f === :mlir_data && !isnothing(v)
-        @assert size(MLIR.IR.type(v)) == ()
-    end
-    return setfield!(x, f, v)
-end
-
 Base.eltype(::Type{TracedRNumber{T}}) where {T} = T
 
 const WrappedTracedRArray{T,N} = WrappedArray{T,N,TracedRArray,TracedRArray{T,N}}
@@ -318,7 +304,7 @@ for (jlop, hloop) in (
     @eval function $(jlop)(
         @nospecialize(lhs::TracedRNumber{T}), @nospecialize(rhs::TracedRNumber{T})
     ) where {T}
-        return TracedRArray{T}(
+        return TracedRNumber{T}(
             (),
             MLIR.IR.result(
                 MLIR.Dialects.stablehlo.$(hloop)(lhs.mlir_data, rhs.mlir_data), 1
@@ -430,7 +416,6 @@ function elem_apply(f, args::Vararg{Any,Nargs}) where {Nargs}
     residx = 1
 
     for a in linear_results
-        @show a
         if has_residx(a)
             path = get_residx(a)
             set!(result, path[2:end], MLIR.IR.result(res, residx))

src/Tracing.jl

@@ -332,8 +332,8 @@ function make_tracer(
         end
         res = if toscalar
             TracedRNumber{T}((path,), nothing)
-        elseif !isnothing(tobatch)
-            TracedRArray{T,length(tobatch)}((path,), prev.mlir_data, tobatch)
+        elseif tobatch !== nothing
+            error("This should not happen...")
         else
             TracedRArray{T,N}((path,), prev.mlir_data, size(prev))
         end
@@ -358,7 +358,9 @@ function make_tracer(
     @nospecialize(prev::TracedRNumber{T}),
     @nospecialize(path),
     mode;
-    kwargs...
+    tobatch=nothing,
+    toscalar=false,
+    kwargs...,
 ) where {T}
     if mode == ConcreteToTraced
         throw("Cannot trace existing trace type")
@@ -374,7 +376,13 @@ function make_tracer(
         if haskey(seen, prev)
             return seen[prev]
         end
-        res = TracedRNumber{T}((path,), prev.mlir_data)
+        res = if toscalar
+            TracedRNumber{T}((path,), nothing)
+        elseif tobatch !== nothing
+            TracedRArray{T,length(tobatch)}((path,), prev.mlir_data, tobatch)
+        else
+            TracedRNumber{T}((path,), prev.mlir_data)
+        end
         seen[prev] = res
         return res
     end