From 6d1de06f2f748acb657cb742e38c709f59da7e64 Mon Sep 17 00:00:00 2001
From: jumerckx <31353884+jumerckx@users.noreply.github.com>
Date: Wed, 21 May 2025 22:24:37 +0200
Subject: [PATCH] allow tobatch !== nothing if shapes match

`error("This should not happen...")` got hit with automatic function
call insertion and this seemed to fix it.
---
 src/Tracing.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Tracing.jl b/src/Tracing.jl
index e110407c05..9d15abe6f6 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -1303,7 +1303,7 @@ Base.@nospecializeinfer function make_tracer(
         end
         res = if toscalar
             TracedRNumber{T}((path,), nothing)
-        elseif tobatch !== nothing
+        elseif tobatch !== nothing && prev.shape != tobatch
             error("This should not happen...")
         else
             TracedRArray{T,N}((path,), prev.mlir_data, size(prev))