Generalize Base._cat to non-Val, typed Base._cat_t and implement typed_hcat, typed_vcat, typed_hvcat, typed_hvncat

src/TracedRArray.jl

@@ -759,10 +759,14 @@ function _copyto!(dest::TracedRArray, bc::Broadcasted)
     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."
+dispatch_val(x) = x
+dispatch_val(::Val{D}) where {D} = D
 
-    # MLIR expects the dimension `D` to be ≤ the rank of the input tensors
+function Base._cat(dims, A::TracedRArray{T,N}, Bs::TracedRArray...) where {T,N}
+    dims = dispatch_val(dims)
+    @assert dims isa Integer "Support for non-integer dimensions is not implemented yet."
+
+    # MLIR expects the dimension `dims` to be ≤ the rank of the input tensors
     A = maybe_expand_dims(A, dims)
     Bs = maybe_expand_dims.(Bs, (dims,))
 
@@ -775,7 +779,7 @@ function Base._cat(dims::Val{D}, A::TracedRArray{T,N}, Bs::TracedRArray...) wher
             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
+                dimension=dims - 1, # stablehlo expects this to be zero-indexed
             ),
             1,
         ),