if (failed(MatchBinaryReduceFunction<BinaryOp>(
              scatter_op.getUpdateComputation()))) {
        return failure();
      }

      auto scatter_dimension_numbers = scatter_op.getScatterDimensionNumbers();

      // Normalize indices so index_vector_dim == indices.rank() - 1.
      int64_t index_vector_dim = scatter_dimension_numbers.getIndexVectorDim();