// CHECK:           return %[[VAL_2]] : tensor<10x8xi64>
// CHECK:         }
func.func @convert_remainder_for_int64(%arg0: tensor<10x8xi64>, %arg1: tensor<10x8xi64>) -> tensor<10x8xi64> {
  %0 = mhlo.remainder %arg0, %arg1 : tensor<10x8xi64>
  func.return %0 : tensor<10x8xi64>
}

// CHECK-LABEL:   func @not_convert_remainder_for_int16(
// CHECK-NOT:       "tf.Mod"

    %0 = stablehlo.constant dense<500> : tensor<100xi64>
    %1 = stablehlo.constant dense<0> : tensor<100xi64>
    %2 = stablehlo.compare  LT, %arg2, %1 : (tensor<100xi64>, tensor<100xi64>) -> tensor<100xi1>
    %3 = stablehlo.add %arg2, %0 : tensor<100xi64>
    %4 = stablehlo.select %2, %3, %arg2 : tensor<100xi1>, tensor<100xi64>
    %5 = stablehlo.reshape %4 : (tensor<100xi64>) -> tensor<100x1xi64>
    %6 = "stablehlo.scatter"(%arg0, %5, %arg3) ({

  // CHECK: "mhlo.all_reduce"
  // CHECK-SAME: channel_handle = #mhlo.channel_handle<handle = 10001, type = 1>
  // CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // CHECK: mhlo.add
  // CHECK: mhlo.return
  %0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Add", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>

  func.func @xla_spmd_shard_to_full_shape(%arg0: tensor<1x2xi64>) -> (tensor<2x2xi64>) {
    // CHECK: %[[SHARDING:.*]] = mhlo.custom_call @Sharding(%arg0) {backend_config = "", mhlo.sharding = "{manual}"} : (tensor<1x2xi64>) -> tensor<1x2xi64>
    // CHECK: %[[FULL:.*]] = mhlo.custom_call @SPMDShardToFullShape(%[[SHARDING]]) {backend_config = "", mhlo.sharding = "{devices=[2,1]0,1}"} : (tensor<1x2xi64>) -> tensor<2x2xi64>
    // CHECK: return %[[FULL]]

  %padding = "tf.Const"() { value = dense<[[1, 2]]> : tensor<1x2xi64> } : () -> tensor<1x2xi64>
  // CHECK: "mhlo.pad"(%arg0, %arg1) <{
  // CHECK-SAME: edge_padding_high = dense<2> : tensor<1xi64>,
  // CHECK-SAME: edge_padding_low = dense<1> : tensor<1xi64>,
  // CHECK-SAME: interior_padding = dense<0> : tensor<1xi64>
  %1 = "tf.PadV2"(%arg0, %padding, %arg1) : (tensor<3xf32>, tensor<1x2xi64>, tensor<f32>) -> tensor<6xf32>

func.func @multinomial_i64(%arg0: tensor<2xf32>, %arg1: tensor<1xi32>) -> tensor<10xi64> {
  %0 = "tf.Multinomial"(%arg0, %arg1) {seed = 0 : i64, seed2 = 0: i64} : (tensor<2xf32>, tensor<1xi32>) -> tensor<10xi64>
  func.return %0 : tensor<10xi64>

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

  %1 = "tf.StringToHashBucketFast"(%arg0) {device = "", num_buckets = 2147483647 : i64} : (tensor<?x!tf_type.string>) -> tensor<?xi64>
  %2 = "tf.Sgnn"(%1, %0) {device = ""} : (tensor<?xi64>, tensor<10x1xi64>) -> tensor<10x?xf64>
  %3 = "tf.Const"() {value = dense<[-1, 10]> : tensor<2xi64>} : () -> tensor<2xi64>

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