[XLA] Fix dynamic expression edge cases and batch bucketing

tensorflow/compiler/jit/BUILD

@@ -2017,4 +2017,14 @@ cc_library(
         "//tensorflow/core/platform:logging",
         "@local_xla//xla:debug_options_flags",
     ],
-)
+)
+
+tf_cc_test(
+    name = "xla_batch_matcher_test",
+    srcs = ["xla_batch_matcher_test.cc"],
+    deps = [
+        ":xla_batch_matcher",
+        "//tensorflow/core:test",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/compiler/jit/xla_batch_matcher.cc

@@ -1,4 +1,7 @@
 #include "tensorflow/compiler/jit/xla_batch_matcher.h"
+
+#include <utility>
+
 #include "xla/debug_options_flags.h"
 #include "tensorflow/core/platform/logging.h"
 
@@ -9,6 +12,14 @@ XlaBatchMatcher::XlaBatchMatcher() {
   parse_env_config();
 }
 
+XlaBatchMatcher::XlaBatchMatcher(std::vector<int64_t> batches)
+    : all_batches_(std::move(batches)) {
+  std::sort(all_batches_.begin(), all_batches_.end());
+  all_batches_.erase(
+      std::unique(all_batches_.begin(), all_batches_.end()),
+      all_batches_.end());
+}
+
 // Trim whitespace (spaces/tabs) from both ends of a string
 std::string trim(const std::string& s) {
   size_t start = s.find_first_not_of(" \t");
@@ -135,7 +146,7 @@ static int64_t GetNextPowerOfTwo(int64_t real_batch) {
   return power;
 }
 
-int64_t XlaBatchMatcher::find_min_larger_batch(int64_t real_batch) {
+int64_t XlaBatchMatcher::find_min_larger_batch(int64_t real_batch) const {
   if (real_batch <= 0 || real_batch > kMaxBatch) {
     LOG(INFO) << "[XLA_BATCH_WARN] Out of valid range: " << real_batch;
     return real_batch;
@@ -152,25 +163,19 @@ int64_t XlaBatchMatcher::find_min_larger_batch(int64_t real_batch) {
   }
   // Edge case 2: Real value ≥ the largest batch, use the nearest power of two
   if (real_batch > all_batches_.back()) {
-    int64_t val = GetNextPowerOfTwo(real_batch);
-    all_batches_.emplace_back(val);
-    print_all_batches();
-    return val;
+    return GetNextPowerOfTwo(real_batch);
   }
 
   // Find first batch larger than real value (binary search via lower_bound)
   auto it = std::lower_bound(all_batches_.begin(), all_batches_.end(), real_batch);
   return (it != all_batches_.end()) ? *it : all_batches_.back();
 }
 
-int64_t XlaBatchMatcher::get_xla_compile_batch(int64_t real_batch) {
+int64_t XlaBatchMatcher::get_xla_compile_batch(int64_t real_batch) const {
   // Match target batch size
   int64_t selected = find_min_larger_batch(real_batch);
-  if (real_batch != last_batch_ || all_batches_.empty()) {
-    last_batch_ = real_batch;
-    VLOG(2) << "[XLA_BATCH_INFO] Real batch: " << real_batch
-      << " -> Selected compile batch: " << selected;
-  }
+  VLOG(2) << "[XLA_BATCH_INFO] Real batch: " << real_batch
+          << " -> Selected compile batch: " << selected;
   return selected;
 }
 

tensorflow/compiler/jit/xla_batch_matcher.h

@@ -20,19 +20,19 @@ constexpr int kMaxBatch = 2147483648ULL >> 1;
 class XlaBatchMatcher {
   public:
     XlaBatchMatcher();
+    explicit XlaBatchMatcher(std::vector<int64_t> batches);
     virtual ~XlaBatchMatcher() = default;
-    int64_t get_xla_compile_batch(int64_t real_batch);
-    std::vector<int64_t> get_all_batches() { return all_batches_; }
+    int64_t get_xla_compile_batch(int64_t real_batch) const;
+    std::vector<int64_t> get_all_batches() const { return all_batches_; }
 
   private:
     void parse_env_config();
     void print_all_batches();
     std::vector<int64_t> parse_single_item(const std::string& item);
-    int64_t find_min_larger_batch(int64_t real_batch);
+    int64_t find_min_larger_batch(int64_t real_batch) const;
 
     std::vector<int64_t> all_batches_;
     std::string env_str_;
-    int64_t last_batch_ = -1;
 };
 
 }  // namespace tensorflow

tensorflow/compiler/jit/xla_batch_matcher_test.cc

@@ -0,0 +1,41 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/jit/xla_batch_matcher.h"
+
+#include <cstdint>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+namespace tensorflow {
+namespace {
+
+TEST(XlaBatchMatcherTest, GeneratedBucketsAreRequestOrderIndependent) {
+  XlaBatchMatcher matcher({100});
+
+  EXPECT_EQ(matcher.get_xla_compile_batch(300), 512);
+  EXPECT_EQ(matcher.get_xla_compile_batch(150), 256);
+}
+
+TEST(XlaBatchMatcherTest, GeneratedBucketsDoNotChangeConfiguredBuckets) {
+  XlaBatchMatcher matcher({200, 100, 200});
+
+  EXPECT_EQ(matcher.get_xla_compile_batch(300), 512);
+  EXPECT_EQ(matcher.get_all_batches(), (std::vector<int64_t>{100, 200}));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/util/strided_slice_op.cc

@@ -211,7 +211,8 @@ absl::Status ValidateStridedSliceOp(
   absl::InlinedVector<xla::DExpr, 4UL> b;
   absl::InlinedVector<xla::DExpr, 4UL> e;
 
-  // HACK
+  // Callers that only request concrete bounds still need matching expressions
+  // for the shared validation logic below.
   if (begin_expr == nullptr) {
     for (int i : *begin) {
       b.push_back(xla::DExpr::Const(i));
@@ -227,7 +228,7 @@ absl::Status ValidateStridedSliceOp(
 
   if (input_shape.unknown_rank()) {
     // Note: If the rank is unknown, "input_shape.dims()" is -1.
-  return errors::InvalidArgument("Unexpected input_shape with unknown rank");
+    return errors::InvalidArgument("Unexpected input_shape with unknown rank");
   }
 
   const bool begin_is_wrong =

third_party/xla/xla/service/cpu/cpu_executable.cc

@@ -327,28 +327,11 @@ absl::Status CpuExecutable::ExecuteComputeFunction(
   return absl::OkStatus();
 }
 
-void PrintScalars(absl::Span<MaybeOwningDeviceMemory const> buffers) {
-  for (int i = 0; i < buffers.size(); ++i) {
-    const se::DeviceMemoryBase& dmem = buffers[i].AsDeviceMemoryBase();
-    if (dmem.opaque() && dmem.size() >= sizeof(int64_t)) {
-      int64_t val = 0;
-      std::memcpy(&val, dmem.opaque(), sizeof(val));
-      std::cerr << "Buffer " << i << " scalar: " << val << std::endl;
-    } else {
-      std::cerr << "Buffer " << i << " empty or too small" << std::endl;
-    }
-  }
-}
-
 absl::Status CpuExecutable::ExecuteThunks(
     const ExecutableRunOptions* run_options,
     absl::Span<MaybeOwningDeviceMemory const> buffers) {
   uint64_t start_ns = tsl::Env::Default()->NowNanos();
 
-  #if defined(PRINT_BATCHSIZE)
-  PrintScalars(buffers);
-  #endif
-
   size_t profile_counters_size = 0;
   int64_t* profile_counters = nullptr;
 

third_party/xla/xla/service/llvm_ir/llvm_util.cc

@@ -886,10 +886,12 @@ static llvm::Value* EmitExpressionImpl(llvm::IRBuilderBase* b,
                                        const DynExpr& expr) {
   llvm::LLVMContext& ctx = b->getContext();
   llvm::IntegerType* i64Type = llvm::IntegerType::getInt64Ty(ctx);
-  if (expr.is_constant()) return llvm::ConstantInt::get(i64Type, expr.get_val(), true);
   if (expr.kind() == DExpr::Kind::kUnknown) {
     return nullptr;
   }
+  if (expr.is_constant()) {
+    return llvm::ConstantInt::get(i64Type, expr.get_val(), true);
+  }
   if (expr.kind() == DExpr::Kind::kVariable) {
     // For now we can just use %bdim...
     return GetBatchDimByName(b);

third_party/xla/xla/service/llvm_ir/llvm_util_test.cc

@@ -54,6 +54,14 @@ TEST(LlvmUtilTest, DynamicExpressionDivisionIsSigned) {
   EXPECT_EQ(division->getOpcode(), llvm::Instruction::SDiv);
 }
 
+TEST(LlvmUtilTest, UnknownExpressionDoesNotEmitAConstant) {
+  llvm::LLVMContext context;
+  llvm::IRBuilder<> builder(context);
+
+  DExpr expression = DExpr::Unknown(kMissingExpressionSentinel);
+  EXPECT_EQ(EmitExpression(&builder, expression), nullptr);
+}
+
 }  // namespace
 }  // namespace llvm_ir
 }  // namespace xla

third_party/xla/xla/shape.cc

@@ -99,9 +99,6 @@ Shape::Shape(const ShapeProto& shape_proto) {
 }
 
 absl::StatusOr<Shape> Shape::FromProto(const ShapeProto& shape_proto) {
-
-  // LOG(INFO) << "FROM PROTO:\n" << shape_proto.DebugString() << std::endl;
-
   Shape shape;
   shape.set_element_type(shape_proto.element_type());
   if (auto* const state = shape.if_array_state()) {
@@ -154,16 +151,13 @@ absl::StatusOr<Shape> Shape::FromProto(const ShapeProto& shape_proto) {
     TF_ASSIGN_OR_RETURN(*shape.mutable_layout(),
                         Layout::FromProto(shape_proto.layout()));
   }
-  // LOG(INFO) << "FROM PROTO " << shape << "\n";
   return shape;
 }
 
 ShapeProto Shape::ToProto() const {
   ShapeProto proto;
   proto.set_element_type(element_type_);
 
-  // LOG(INFO) << "TO PROTO " << ToString() << "\n";
-
   if (const auto* const state = if_array_state()) {
     proto.mutable_dimensions()->Reserve(state->dimensions.size());
     for (const int64_t dimension : state->dimensions) {
@@ -189,7 +183,6 @@ ShapeProto Shape::ToProto() const {
     proto.mutable_tuple_shapes()->Reserve(1);
     *proto.add_tuple_shapes() = state->buffer_shape[0].ToProto();
   }
-  // LOG(INFO) << "DEBUG VIEW:\n" << proto.DebugString() << std::endl;
   return proto;
 }
 

third_party/xla/xla/shape_expr.cc

@@ -468,6 +468,9 @@ std::unique_ptr<DynExpr> SimplifyFallback(const DynExpr* expr) {
       }
       Constant* l = AsConstant(lhs.get());
       Constant* r = AsConstant(rhs.get());
+      if (r && r->get_val() == 0) {
+        return std::make_unique<Div>(lhs.release(), rhs.release());
+      }
       if (*lhs == *rhs) {
         return std::make_unique<Constant>(1);
       }

third_party/xla/xla/shape_expr.h

@@ -66,7 +66,9 @@ class DynExpr {
   virtual std::set<int> get_all_ids() = 0;
   virtual std::optional<int64_t> solve(int64_t x) = 0;
 
-  bool is_dynamic() { return !is_constant(); }
+  bool is_dynamic() const {
+    return kind() != DExprKind::kUnknown && !is_constant();
+  }
 
   static DynExpr* zero;
   static DynExpr* one;
@@ -193,7 +195,7 @@ class UnknownExpr : public DynExpr {
   void to_proto(xla::ExpressionProto* proto) const override {
     (void)proto;
   }
-  bool is_constant() const override { return true; }
+  bool is_constant() const override { return false; }
   int get_id() const { return id_; }
   DynExpr* substitute(int id, DynExpr* v) override {
     (void)id;

third_party/xla/xla/shape_test.cc

@@ -133,6 +133,21 @@ TEST_F(ShapeTest, DExprSolveRejectsZeroDivisor) {
   EXPECT_FALSE(expr->solve(7).has_value());
 }
 
+TEST_F(ShapeTest, DExprSimplifyPreservesZeroDivisor) {
+  DExpr expr = DExpr::Const(0) / DExpr::Const(0);
+  DExpr simplified = expr.simplify();
+
+  EXPECT_EQ(DExpr::Kind::kDiv, simplified.kind());
+  EXPECT_FALSE(simplified->is_constant());
+}
+
+TEST_F(ShapeTest, UnknownExpressionIsNeitherConstantNorDynamic) {
+  DExpr expr = DExpr::Unknown(kMissingExpressionSentinel);
+
+  EXPECT_FALSE(expr->is_constant());
+  EXPECT_FALSE(expr->is_dynamic());
+}
+
 TEST_F(ShapeTest, DExprMaxSimplifiesAndRoundTrips) {
   DExpr expr = DExpr::Max(DExpr::Var(1), DExpr::Const(4));
   EXPECT_EQ("max(A, 4)", DExprToString(expr.simplify()));