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

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

func.func @stridedSliceGrad(%dy: tensor<4x8xf32>, %begin: tensor<2xi64>, %end: tensor<2xi64>, %strides: tensor<2xi64>, %shape: tensor<1x2xi64>) -> tensor<?x?xf32> {
  // expected-error @+1 {{'shape' operand must be 1D tensor, but got 2D tensor}}
  %0 = "tf.StridedSliceGrad"(%shape, %begin, %end, %strides, %dy) : (tensor<1x2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<4x8xf32>) -> tensor<?x?xf32>
  func.return %0 : tensor<?x?xf32>
}