func.return %0 : tensor<3x4xi1>
}

// CHECK-LABEL: @is_nan
func.func @is_nan(%arg0: tensor<3x4xf32>) -> tensor<3x4xi1> {
  // CHECK: %[[RESULT:.*]] = "tf.NotEqual"(%arg0, %arg0) <{incompatible_shape_error = true}> : (tensor<3x4xf32>, tensor<3x4xf32>) -> tensor<3x4xi1>
  %0 = "tf.IsNan"(%arg0) : (tensor<3x4xf32>) -> tensor<3x4xi1>
  // CHECK: return %[[RESULT]]
  func.return %0 : tensor<3x4xi1>
}

func.func @broadcast_to_i16_low_dim(%input: tensor<3xi16>, %shape: tensor<2xi32>) -> tensor<3x3xi16> {
  %0 = "tf.BroadcastTo"(%input, %shape) : (tensor<3xi16>, tensor<2xi32>) -> tensor<3x3xi16>
  func.return %0: tensor<3x3xi16>

// CHECK-LABEL: broadcast_to_i16_low_dim
// CHECK:    %0 = "tf.BroadcastTo"(%arg0, %arg1) : (tensor<3xi16>, tensor<2xi32>) -> tensor<3x3xi16>
// CHECK:    return %0 : tensor<3x3xi16>
}

// CHECK-LABEL: QuantizePad
func.func @QuantizePad(tensor<2x1x3x!quant.uniform<u8:f32, 0.1>>, tensor<3x2xi32>) -> tensor<?xf32> {
^bb0(%arg0: tensor<2x1x3x!quant.uniform<u8:f32, 0.1>>, %arg1: tensor<3x2xi32>):
  %0 = "tfl.dequantize"(%arg0) : (tensor<2x1x3x!quant.uniform<u8:f32, 0.1>>) -> tensor<2x1x3xf32>
  %1 = "tfl.pad"(%0, %arg1) : (tensor<2x1x3xf32>, tensor<3x2xi32>) -> tensor<?xf32>
  func.return %1 : tensor<?xf32>

// CHECK: %[[PAD:.*]] = "tf.PadV2"({{.*}}, %[[CONST]], %[[CONST_1]])
// CHECK: %[[CONV:.*]] = "tf.XlaConvV2"(%[[PAD]], %[[WEIGHT]]
// CHECK-SAME: (tensor<1x4x5x5x3xi8>, tensor<2x3x3x3x2xi8>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor<i32>) -> tensor<1x3x2x3x2xi32>
// CHECK: %[[SUB:.*]] = "tf.Sub"(%[[CONV]], %[[CONST_2]])
}

// -----

module attributes {tf_saved_model.semantics} {