// CHECK: "tf.Conv3DBackpropInputV2"
}

func.func @mul_i64(%arg0: tensor<14xi64>, %arg1: tensor<14xi64>) -> tensor<14xi64> {
  %0 = "tf.Mul"(%arg0, %arg1) : (tensor<14xi64>, tensor<14xi64>) -> tensor<14xi64>
  func.return %0: tensor<14xi64>

// CHECK-LABEL: mul_i64
// CHECK:  tfl.mul %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<14xi64>
// CHECK:  return
}

  %begins = "tf.Const"() {value = dense<[-1]> : tensor<1xi64>} : () -> (tensor<1xi64>)
  %sizes = "tf.Const"() {value = dense<[2]> : tensor<1xi64>} : () -> (tensor<1xi64>)
  // expected-error @+1 {{requires 0 <= begin[i] <= begin[i] + size[i] <= Di}}
  %0 = "tf.Slice"(%arg0, %begins, %sizes) : (tensor<4xi32>, tensor<1xi64>, tensor<1xi64>) -> tensor<2xi32>
  func.return %0 : tensor<2xi32>
}

// -----

// 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>
// CHECK-DAG:       [[VAL_39:%.*]] = arith.constant dense<1> : tensor<1xi64>

// CHECK-DAG: %[[ZERO:.+]] = arith.constant dense<0> : tensor<1xi64>
// CHECK-DAG: %[[MAX1:.+]] = arith.constant dense<2> : tensor<1xi64>
// CHECK-DAG: %[[MAX2:.+]] = arith.constant dense<3> : tensor<1xi64>
// CHECK: %[[BITCAST1:.+]] = "tfl.bitcast"(%[[ARG1]]) : (tensor<i64>) -> tensor<1xi64>
// CHECK: %[[MIN1:.+]] = "tfl.minimum"(%[[BITCAST1]], %[[MAX1]]) : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi64>

    // CHECK:             %[[A_SHARD:.+]] = "tf.XlaSpmdFullToShardShape"(%[[A]]) <{dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []}> : (tensor<2x2xi64>) -> tensor<1x2xi64>
    // CHECK:             %[[B:.+]] = "tf._XlaHostComputeMlir"(%[[A_SHARD]])
    // CHECK-SAME:          manual_sharding = true
    // CHECK-SAME:          recv_key = "host_compute_channel_0_retvals"

    func.return %0 : tensor<*xf32>
  }

  // Verifies handling of cases involving multiple iteration of feeding inputs.
  // CHECK-LABEL: @const_input_required
  func.func @const_input_required(%arg0: tensor<10xf64>) -> tensor<?xf64> attributes {tf.entry_function = {control_outputs = "", inputs = "_arg0,_arg1,_arg2,_arg3", outputs = "_retval0"}} {
    %cst = "tf.Const"() {value = dense<6> : tensor<1xi32>} : () -> tensor<1xi32>

// CHECK-LABEL: testMultinomial
func.func @testMultinomial(%arg0: tensor<2xf32>, %arg1: tensor<1xi32>) -> tensor<10xi64> {
  // CHECK: "tfl.multinomial"(%arg0, %arg1)
  %0 = "tfl.multinomial"(%arg0, %arg1) {seed = 0 : i64, seed2 = 0: i64} : (tensor<2xf32>, tensor<1xi32>) -> tensor<10xi64>
  func.return %0 :  tensor<10xi64>
}

// -----

// CHECK-LABEL: testMultinomialInt32

    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 =