}

func.func @broadcast_args_i64(%arg0: tensor<3xi64>, %arg1: tensor<1xi64>) -> tensor<3xi64> {
  %0 = "tf.BroadcastArgs"(%arg0, %arg1) : (tensor<3xi64>, tensor<1xi64>) -> tensor<3xi64>
  func.return %0 : tensor<3xi64>

// CHECK-LABEL:broadcast_args_i64
// CHECK:  "tfl.broadcast_args"(%arg0, %arg1) : (tensor<3xi64>, tensor<1xi64>) -> tensor<3xi64>
}

  %0 = "tfl.shape"(%arg1) : (tensor<8x7x6x5x?x3x2x1xf32>) -> tensor<8xi64>
  %1 = "tfl.shape"(%arg2) : (tensor<?x3x2x1xf32>) -> tensor<4xi64>
  %2 = "tfl.broadcast_args"(%0, %1) : (tensor<8xi64>, tensor<4xi64>) -> tensor<8xi64>
  %3 = "tfl.shape"(%arg0) : (tensor<8x7x6x5x?x3x2x1xi1>) -> tensor<8xi64>
  %4 = "tfl.broadcast_args"(%2, %3) : (tensor<8xi64>, tensor<8xi64>) -> tensor<8xi64>

  %end = "tf.Const"() { value = dense<[5, 10]> : tensor<2xi64> } : () -> tensor<?xi64>
  %strides = "tf.Const"() { value = dense<[2, 3, 4]> : tensor<3xi64> } : () -> tensor<?xi64>

  // expected-error @+1 {{to have the same number of elements}}
  %1 = "tf.StridedSlice"(%input, %begin, %end, %strides) : (tensor<4x8xf32>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>) -> tensor<?x?xf32>
}

// CHECK-DAG:       [[VAL_35:%.*]] = arith.constant dense<0> : tensor<1xi64>
// CHECK-DAG:       [[VAL_36:%.*]] = arith.constant dense<1> : tensor<1xi64>
// CHECK:           [[VAL_37:%.*]] = "tf.Slice"([[VAL_3]], [[VAL_35]], [[VAL_36]]) : (tensor<2xf32>, tensor<1xi64>, tensor<1xi64>) -> tensor<1xf32>
// CHECK-DAG:       [[VAL_38:%.*]] = arith.constant dense<1> : tensor<1xi64>

  func.return %3 : tensor<2xi64>
  // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi64>}> : () -> tensor<2xi64>
  // CHECK: return %[[CST]]
}

// CHECK-LABEL: testFoldStridedSliceShapeDynamicOutput
func.func @testFoldStridedSliceShapeDynamicOutput(%arg0: tensor<?x1x2x?xf32>) -> (tensor<?xi32>) {

  func.func @map_outside_compilation_0_inputs(%arg0 : tensor<2xi64>) -> () {
    // CHECK:             "tf._XlaHostComputeMlir"()
    "tf_device.cluster"() ({
      %1 = "tf.OpA"(%arg0) {_xla_map_outside_compilation = "0", _xla_outside_compilation = "from_launch"} : (tensor<2xi64>) -> tensor<2xi64>
      "tf.OpB"(%1) : (tensor<2xi64>) -> ()
      tf_device.return

// CHECK: %[[BITCAST2:.+]] = "tfl.bitcast"(%[[ARG2]]) : (tensor<i64>) -> tensor<1xi64>
// CHECK: %[[MIN2:.+]] = "tfl.minimum"(%[[BITCAST2]], %[[MAX2]]) : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi64>
// CHECK: %[[CONCAT:.+]] = "tfl.concatenation"(%[[MIN1]], %[[MIN2]]) <{axis = 0 : i32, fused_activation_function = "NONE"}> : (tensor<1xi64>, tensor<1xi64>) -> tensor<2xi64>

    %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>

    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,