From 176aa7bc013ffc33b2e001ae3c8c81a3042cf06f Mon Sep 17 00:00:00 2001
From: Misha Gutman <aelphy@google.com>
Date: Fri, 13 Mar 2026 09:49:32 -0700
Subject: [PATCH] Added non trivial batch dim.

PiperOrigin-RevId: 883212002
---
 bench/subgraph/static-reduce.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/bench/subgraph/static-reduce.cc b/bench/subgraph/static-reduce.cc
index b2d91b90def..829e96e3cd2 100644
--- a/bench/subgraph/static-reduce.cc
+++ b/bench/subgraph/static-reduce.cc
@@ -77,12 +77,12 @@ static void FP32Reduce(benchmark::State& state) {
   const size_t d1 = state.range(2);
   const size_t d2 = state.range(3);
   const size_t norm_mask = state.range(4);
-  std::vector<size_t> dims = {d0, d1, d2};
+  std::vector<size_t> dims = {1000, d0, d1, d2};
   std::vector<size_t> axes;
   axes.reserve(dims.size());
   for (size_t i = 0; i < dims.size(); ++i) {
     if ((norm_mask & (1 << i)) != 0) {
-      axes.push_back(i);
+      axes.push_back(i + 1);
     }
   }