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();

  if (subgraph_index < 0) return std::nullopt;

  mlir::stablehlo::ScatterDimensionNumbersAttr scatter_dimension_numbers =
      scatter_op.getScatterDimensionNumbers();
  llvm::ArrayRef<int64_t> update_window_dims_mlir =
      scatter_dimension_numbers.getUpdateWindowDims();
  llvm::ArrayRef<int64_t> inserted_window_dims_mlir =