refactor: move code a bit

Author: avik-pal

src/TracedRArray.jl

@@ -455,18 +455,41 @@ function Base.fill!(A::TracedRArray{T,N}, x::TracedRNumber{T2}) where {T,N,T2}
     return A
 end
 
+function Base._cat(dims::Val{D}, A::TracedRArray{T,N}, Bs::TracedRArray...) where {T,N,D}
+    @assert D isa Integer "Support for non-integer dimensions is not implemented yet."
+
+    # MLIR expects the dimension `D` to be ≤ the rank of the input tensors
+    A = maybe_expand_dims(A, dims)
+    Bs = maybe_expand_dims.(Bs, (dims,))
+
+    catdims = Base.dims2cat(dims)
+    shape = Base.cat_size_shape(catdims, A, Bs...)
+    RT = Base.promote_eltype(A, Bs...)
+    Res = TracedRArray{RT,length(shape)}(
+        (),
+        MLIR.IR.result(
+            MLIR.Dialects.stablehlo.concatenate(
+                [A.mlir_data, [B.mlir_data for B in Bs]...];
+                result_0=MLIR.IR.TensorType(shape, MLIR.IR.Type(RT)),
+                dimension=D - 1, # stablehlo expects this to be zero-indexed
+            ),
+            1,
+        ),
+        shape,
+    )
+    return Res
+end
+
+function maybe_expand_dims(x::AbstractArray{T,N}, ::Val{D}) where {T,N,D}
+    D ≤ N && return x
+    return reshape(x, ntuple(i -> i ≤ N ? size(x, i) : 1, Val(D)))
+end
+
 struct AbstractReactantArrayStyle{N} <: Base.Broadcast.AbstractArrayStyle{N} end
 
 AbstractReactantArrayStyle(::Val{N}) where {N} = AbstractReactantArrayStyle{N}()
 AbstractReactantArrayStyle{M}(::Val{N}) where {N,M} = AbstractReactantArrayStyle{N}()
 
-# function Broadcast.materialize(bc::Broadcasted) 
-#    @show bc
-#    inst = instantiate(bc)
-#    @show inst
-#    copy(inst)
-# end
-
 function BroadcastStyle(::Type{<:AnyTracedRArray{T,N}}) where {T,N}
     return AbstractReactantArrayStyle{N}()
 end
@@ -628,33 +651,3 @@ function _copyto!(dest::TracedRArray, bc::Broadcasted)
     dest.mlir_data = res.mlir_data
     return dest
 end
-
-function Base._cat(dims::Val{D}, A::TracedRArray{T,N}, Bs::TracedRArray...) where {T,N,D}
-    @assert D isa Integer "Support for non-integer dimensions is not implemented yet."
-
-    # MLIR expects the dimension `D` to be ≤ the rank of the input tensors
-    A = maybe_expand_dims(A, dims)
-    Bs = maybe_expand_dims.(Bs, (dims,))
-
-    catdims = Base.dims2cat(dims)
-    shape = Base.cat_size_shape(catdims, A, Bs...)
-    RT = Base.promote_eltype(A, Bs...)
-    Res = TracedRArray{RT,length(shape)}(
-        (),
-        MLIR.IR.result(
-            MLIR.Dialects.stablehlo.concatenate(
-                [A.mlir_data, [B.mlir_data for B in Bs]...];
-                result_0=MLIR.IR.TensorType(shape, MLIR.IR.Type(RT)),
-                dimension=D - 1, # stablehlo expects this to be zero-indexed
-            ),
-            1,
-        ),
-        shape,
-    )
-    return Res
-end
-
-function maybe_expand_dims(x::AbstractArray{T,N}, ::Val{D}) where {T,N,D}
-    D ≤ N && return x
-    return reshape(x, ntuple(i -> i ≤ N ? size(x, i) : 1, Val(D)))
-end