%0 = "tf.BroadcastTo"(%arg0, %arg1) : (tensor<3xbf16>, tensor<2xi64>) -> tensor<3x3xbf16>
  func.return %0: tensor<3x3xbf16>

// CHECK-LABEL: broadcast_to_bf16
// CHECK:  [[CST:%.*]] = arith.constant dense<1.000000e+00> : tensor<3x3xbf16>
// CHECK:  [[MUL:%.*]] = tfl.mul(%arg0, [[CST]]) <{fused_activation_function = "NONE"}> : (tensor<3xbf16>, tensor<3x3xbf16>) -> tensor<3x3xbf16>

// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -emit-custom-ops -o - | flatbuffer_to_string - | FileCheck %s

func.func @main(tensor<4x5xbf16>) -> tensor<4x5xbf16> {
^bb0(%arg0: tensor<4x5xbf16>):

// CHECK:  {
// CHECK-NEXT:      version: 3,
// CHECK-NEXT:      operator_codes: [ {
// CHECK-NEXT:        deprecated_builtin_code: 53,
// CHECK-NEXT:        version: 7,
// CHECK-NEXT:        builtin_code: CAST
// CHECK-NEXT:      } ],

  %0 = "tf.Dequantize"(%arg0, %min_range, %max_range) : (tensor<2x3x!tf_type.qint8>, tensor<f32>, tensor<f32>) -> tensor<2x3xbf16>

  // CHECK-DAG: return %[[RESULT]]
  func.return %0 : tensor<2x3xbf16>

    %1 = "tf.Cast"(%0) <{truncate = false}> : (tensor<3x1xf32>) -> tensor<3x1xbf16>
    %2 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
    "tf.AssignVariableOp"(%2, %1) : (tensor<!tf_type.resource<tensor<3x1xbf16>>>, tensor<3x1xbf16>) -> ()
    return
  }

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

func.func @sparse_qu8() -> tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.0>> {
  // CHECK-LABEL: @sparse_qu8

      shape {
        unknown_rank: true
      }
    }
  }
  attr {
    key: "dtype"
    value {
      type: DT_BOOL
    }
  }
}

# CHECK-DAG: tf.IteratorGetNext{{.+}}-> (tensor<1x8xbf16>, tensor<2x?x16xf32>, tensor<*xf64>)
# CHECK-DAG: tf.IteratorGetNextSync{{.+}}-> (tensor<*xi16>, tensor<3x24xi32>, tensor<?x4x32xi64>)

      return %0 : tensor<1x?x4xbf16>
    }
  })";

  TF_ASSERT_OK(LegalizeModule(kModuleWithNonConstParam));
}

TEST_F(Tf2XlaRewriterTest, ErrorsWithInvalidNumberOfParametersToArgs) {

        tensor_shape {
          dim {
            size: 2
          }
        }
        half_val: 16964
        half_val: 17485
        # CHECK: value = dense<[4.900000e+01, 8.200000e+02]> : tensor<2xbf16>
      }
    }
  }
}
node {
  name: "double"
  op: "Const"
  device: "/device:CPU:0"
  attr {
    key: "dtype"
    value {
      type: DT_DOUBLE
    }
  }
  attr {
    key: "value"

// Test valid tf.Select
// CHECK-LABEL: func @testSelect
func.func @testSelect(%arg0: tensor<3xi1>, %arg1: tensor<3x2xf16>, %arg2: tensor<3x2xf16>) -> tensor<3x2xf16> {
  %0 = "tf.Select"(%arg0, %arg1, %arg2) : (tensor<3xi1>, tensor<3x2xf16>, tensor<3x2xf16>) -> tensor<3x2xf16>
  func.return %0: tensor<3x2xf16>
}

// -----

func.func @floordiv_bf16(%arg0: tensor<2xbf16>) -> tensor<2xbf16> {
  // CHECK-NEXT:  mhlo.convert
  // CHECK-NEXT:  mhlo.convert
  // CHECK-NEXT:  chlo.broadcast_divide
  // CHECK-NEXT:  mhlo.floor
  // CHECK-NEXT:  mhlo.convert
  // CHECK-NEXT:  return
  %0 = "tf.FloorDiv"(%arg0, %arg0) : (tensor<2xbf16>, tensor<2xbf16>) -> tensor<2xbf16>
  func.return %0: tensor<2xbf16>
}

// -----