assert(index[i] >= 0);
    }
#endif  // NDEBUG
    int64_t offset = 0;
    int64_t stride = 1;
    for (int64_t dim = shape.size() - 1; dim >= 0; --dim) {
      offset += index[dim] * stride;
      stride *= shape[dim];
    }
    return offset;
  }

  int64_t offset_;
  int64_t stride_;
  int64_t size_;
};

template <typename T>

  // CHECK: %[[CST_0:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}>
  // CHECK: "tf.Transpose"(%[[R1]], %[[CST_0]])
  %0 = "tf.MaxPool"(%arg0)
       {
         data_format = "NCHW",
         ksize = [1, 1, 3, 3],
         padding = "SAME",
         strides = [1, 1, 2, 2]

  // 1) nullptr indicates tensor is compact and row-majored.
  // 2) fill in the strides array as the real case for compact row-major data.
  // Here we choose option 2, since some frameworks didn't handle the strides
  // argument properly.
  dlm_tensor->dl_tensor.strides = stride_arr->data();

  dlm_tensor->dl_tensor.byte_offset =
      0;  // TF doesn't handle the strides and byte_offsets here
  return static_cast<void*>(dlm_tensor);
}

  GetI32VectorFromDenseI64CompositeAttr(composite_attrs, "padding", &padding);
  data.ph = padding[0];
  data.pw = padding[1];

  std::vector<int32_t> strides;
  GetI32VectorFromDenseI64CompositeAttr(composite_attrs, "stride", &strides);
  data.sh = strides[0];
  data.sw = strides[1];

  data.ceil_mode =
      GetBoolFromCompositeAttr(composite_attrs, "ceil_mode").value();

  if (strides[0] != 1 || strides[3] != 1) return 1;

  // Space to depth only supports height_stride == width_stride case.
  if (strides[1] != strides[2]) return 1;

  return strides[1];
}

void TPUSpaceToDepthPass::runOnOperation() {
  std::optional<tf_device::ClusterFuncOp> cluster_func;
  // Space to depth only supports training loop.

  // CHECK-SAME: data_format = "NCHW"
  // CHECK-SAME: dilations = [1, 4, 2, 3]
  // CHECK-SAME: explicit_paddings = [1, 2, 7, 8, 3, 4, 5, 6]
  // CHECK-SAME: padding = "EXPLICIT"
  // CHECK-SAME: strides = [5, 8, 6, 7]
  // CHECK-SAME: (tensor<1x3x32x32xf32>, tensor<1x1x3x8xf32>) -> tensor<1x8x7x7xf32>

  // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}>

      return failure();

    // We also need constant begin/end indices and strides to perform padding
    // calculations.
    // Bounded shape after performing strided slice
    SmallVector<int64_t, 4> shape;
    // Bounded begin, end, and strides for strided slice
    SmallVector<int64_t, 4> begin_indices, end_indices, strides;
    if (!op.GetSlicedShapeAndBoundRanges(&shape, &begin_indices, &end_indices,

    // CHECK-NEXT: %3 = mhlo.reshape %2 : (tensor<1xi32>) -> tensor<i32>
    // CHECK-NEXT: %4 = mhlo.convert %3 : tensor<i32>
    // CHECK-NEXT: %5 = "mhlo.slice"(%1) <{limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1xi32>

                                                           TFLOp op) {
  TfLitePoolParams pool_params;
  pool_params.stride_height = op.stride_h().getSExtValue();
  pool_params.stride_width = op.stride_w().getSExtValue();
  pool_params.filter_height = op.filter_h().getSExtValue();
  pool_params.filter_width = op.filter_w().getSExtValue();
  const auto padding = GetTflitePadding(inst, op.padding());

  // Dynamic stride: when `begin` and `end` inputs are unknown at compile time,
  // `strides` must be known.
  // CHECK: tf.StridedSlice