%strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi64>} : () -> tensor<3xi64>
    // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) <{dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""}> : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi64>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor<i32>) -> tensor<8x4x14x14x16xf32>

    "tf._XlaHostComputeMlir"() {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", host_mlir_module = ""} : () -> ()
    "mhlo.return"(%1) : (tensor<f32>) -> ()
  }) {dimensions = dense<[1]> : tensor<1xi64>} : (tensor<?x?xf32>, tensor<f32>) -> tensor<?xf32>
  func.return
}

// -----

// Tests transitive unsupported parent of TF/XLA communication op.

//     -> tensor<5x2xf32>
//
// is lowered to
//
//   %shape = "tf.Const"() {value = dense<[-1, 2]> : tensor<2xi64>}
//   %inp0 = "tf.Reshape"(%arg0, %shape)
//     : (tensor<2xf32>, tensor<2xi64>) -> tensor<1x2xf32>
//   %inp1 = "tf.Reshape"(%arg1, %shape)
//     : (tensor<2x2x2xf32>, tensor<2xi64>) -> tensor<4x2xf32>
//   %items0 = "tf.Unpack"(%[[INP0]]) {axis = 0 : i64}
//     : (tensor<1x2xf32>) -> tensor<2xf32>

      while (auto cast_op =
                 llvm::dyn_cast_or_null<CastOp>(value.getDefiningOp())) {
        // Consider cast compatibility in case
        //    %cast = "tf.Cast"(%0) : (tensor<2xi64>) -> tensor<2xf32>
        // is skipped.
        if (cast_op.getSrcT() != cast_op.getDstT()) {
          break;
        }
        value = cast_op.getOperand();
      }
      return value;
    };

    SmallVector<Value> start_indices(rank);
    for (auto [i, start_index] : llvm::enumerate(op.getStartIndices())) {
      // Start indices should be casted from tensor<i64> to tensor<1xi64>.
      auto cast = rewriter.create<TFL::BitcastOp>(
          op->getLoc(), single_element_type, start_index);
      int64_t upper_limit_idx = operand_shape[i] - slice_sizes[i];
      auto upper_limit_attr =

// For example,
//
//   %0 = tf.Pack(%input) {axis = 0} // %input : tensor<2x3xf32>
//
// can be canonicalized to
//
//   %shape = "tf.Const"() {value = dense<[1, 2, 3]> : tensor<3xi64>}
//   %0 = tf.Reshape(%input, %shape)
struct ConvertPackToReshape : public OpRewritePattern<PackOp> {
  using OpRewritePattern<PackOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(PackOp pack_op,