Resolve comments

openxla · Oct 8, 2024 · 9d8eb96 · 9d8eb96
1 parent 82f2237
commit 9d8eb96
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 8 deletions.
diff --git a/xla/service/scatter_determinism_expander_test.cc b/xla/service/scatter_determinism_expander_test.cc
@@ -230,8 +230,7 @@ TEST_F(ScatterDeterminismExpanderTest, ScatterAddHloVerificationTest) {
   )";
 
   RunAndFilecheckHloRewrite(kModuleStr, ScatterDeterminismExpander(),
-                            kExpectedPattern, nullptr /*after_pass_checks*/,
-                            nullptr /*config*/);
+                            kExpectedPattern, nullptr, nullptr);
 }
 
 TEST_F(ScatterDeterminismExpanderTest, ScatterAddOutOfBoundCorrectnessTest) {

diff --git a/xla/service/scatter_utils.cc b/xla/service/scatter_utils.cc
@@ -9,12 +9,7 @@ namespace xla {
 StatusOr<HloInstruction*> TransposeIndexVectorDimToLast(
     HloInstruction* scatter_indices, int64_t index_vector_dim) {
   const Shape& scatter_indices_shape = scatter_indices->shape();
-
-  if (scatter_indices_shape.dimensions_size() == index_vector_dim) {
-    return scatter_indices;
-  }
-
-  if (index_vector_dim == (scatter_indices_shape.dimensions_size() - 1)) {
+  if (index_vector_dim >= (scatter_indices_shape.dimensions_size() - 1)) {
     return scatter_indices;
   }