From bb1221179b214947a3b23fc4a3648a1b2717dba8 Mon Sep 17 00:00:00 2001
From: Pavel Potapov <ppotapov@amazon.com>
Date: Fri, 16 Jan 2026 15:48:37 -0500
Subject: [PATCH] Fix the serde bug for arange

---
 KLR/Trace/Tensor.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/KLR/Trace/Tensor.lean b/KLR/Trace/Tensor.lean
index aad14413..74940b9e 100644
--- a/KLR/Trace/Tensor.lean
+++ b/KLR/Trace/Tensor.lean
@@ -76,7 +76,7 @@ nki builtin.meta.tensor.arange
  (dtype : Dtype := .float32) := do
   let tlDtype <- dtype.toTensorLibDtype
   let cnt := (stop - start + step - 1).div  step
-  let tlShape := TensorLib.Shape.mk [cnt]
+  let tlShape := TensorLib.Shape.mk [1, cnt]
   let values := List.range cnt |>.map (fun i => start + i * step)
   let mut data := ByteArray.emptyWithCapacity $ cnt * tlDtype.itemsize
   for v in values do