// should have the same size.
// Example: Index (2, 3) of a (4, 5)-shaped tensor has the contiguous offset of
// 2 * 5 + 3 = 13.
int64_t GetContiguousOffset(const ArrayRef<int64_t> indices,
                            const ArrayRef<int64_t> shape) {
  int64_t contiguous_offset = 0;
  int64_t base_offset = 1;
  for (auto [i, dimension] : llvm::reverse(llvm::zip_equal(indices, shape))) {
    contiguous_offset += base_offset * i;

}

Value CreateStridedSliceOp(mlir::Location loc, ArrayRef<int64_t> output_shape,
                           Value input, ArrayRef<int32_t> begin,
                           ArrayRef<int32_t> end, ArrayRef<int32_t> strides,
                           int64_t begin_mask, int64_t end_mask,
                           int64_t ellipsis_mask, int64_t new_axis_mask,
                           int64_t shrink_axis_mask, OpBuilder* builder) {

// Inversely permutate a given vector
static SmallVector<int64_t> InvertPermutationToVector(ArrayRef<int64_t> vec,
                                                      ArrayRef<int64_t> perm) {
  return applyPermutation(vec, invertPermutationVector(perm));
}

static RankedTensorType GetPermutedTensorTypeHelper(RankedTensorType type,
                                                    ArrayRef<int64_t> perm,

// calculate padding_low and padding_high for SAME padding.
void GetSamePaddingValues(OpBuilder &builder, Location loc, Value input_size,
                          int64_t filter_sz, int64_t dilation_rate,
                          int64_t stride, Value &padding_low,
                          Value &padding_high) {
  Value zero = CreateScalarConstValue<int32_t>(builder, loc, 0);

  TapeVSpace vspace(ctx);
  std::vector<int64_t> target_tensor_ids = MakeTensorIDList(targets);
  std::vector<int64_t> source_tensor_ids = MakeTensorIDList(sources);
  tensorflow::gtl::FlatSet<int64_t> sources_set(source_tensor_ids.begin(),
                                                source_tensor_ids.end());
  std::unordered_map<int64_t, TapeTensor> sources_that_are_targets;

  // Map from tensor ID to how many references still exist for this tensor in
  // the tape.
  std::unordered_map<int64_t, int64_t> tensor_usage_counts;

  // Maps from op ID to how many output tensors of this op still need to have
  // their gradients computed.
  std::unordered_map<int64_t, int64_t> op_missing_tensor;
};

// If `persistent_tape` is true, op_tape is not changed and none of the

  builder.setInsertionPoint(conv2d);
  // Book keeping filter information.
  auto filter_shape = mlir::cast<RankedTensorType>(filter.getType()).getShape();
  int64_t height = filter_shape[0];
  int64_t width = filter_shape[1];
  int64_t channel = filter_shape[2];
  int64_t out_channel = filter_shape[3];
  // Value/Op before reshape op.
  Value before_reshape_value = filter;
  if (height % block_size != 0 || width % block_size != 0) {

    // Otherwise, output type has same rank as the input.
    return RankedTensorType::get(
        SmallVector<int64_t, 4>(rank, ShapedType::kDynamic), element_ty);
  }

  int64_t num_reduce_dim = 0;
  llvm::SmallVector<bool, 4> is_reduce_dim(rank, false);
  for (const APInt &index : indices.getValues<APInt>()) {
    int64_t dim = GetDimForAxis(index.getSExtValue(), rank);
    // Invalid input.

std::optional<std::pair<SmallVector<int64_t>, SmallVector<int64_t>>>
GetBroadcastShapesForBatchMatmul(ShapedType input_type,
                                 ShapedType weight_type) {
  ArrayRef<int64_t> input_shape = input_type.getShape();
  ArrayRef<int64_t> weight_shape = weight_type.getShape();

  const int64_t num_matmul_dim = 2;
  const int64_t num_input_batch_dim = input_type.getRank() - num_matmul_dim;

    auto can_convert_to_bilinear = [](bool align_corners, int64_t dilation,
                                      int64_t padding, int64_t stride,
                                      int64_t input_spatial,
                                      int64_t output_spatial) {
      int64_t input_spatial_size =
          align_corners ? input_spatial - 1 : input_spatial;
      int64_t output_spatial_size =