}

// Decompose tf.RngReadAndSkip.
//
// For Philox, the resource variable holds a tensor<3xi64> with the state:
//   [counter_lo, counter_hi, key]
//
//   RngReadAndSkip increments the 128 bit counter value by 256 * delta and
//   returns the original state value.
//
// For Threefry, the resource variable holds a tensor<2xi64> with the state:
//   [counter, key]
//

    // CHECK: %[[CONST3:.*]] = "tf.Const"() <{value = dense<
    // CHECK-SAME: [7, 7, 3, 64]
    // CHECK: %[[CONST4:.*]] = "tf.Const"() <{value = dense<
    // CHECK-SAME: 0
    // CHECK: %[[SLICE0:.*]] = "tf.Slice"(%[[RESHAPE1:.*]], %[[CONST4:.*]], %[[CONST3:.*]]) : (tensor<8x8x3x64xf32>, tensor<4xi64>, tensor<4xi32>) -> tensor<7x7x3x64xf32>

  func.return %3 : tensor<2xi64>
  // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi64>}> : () -> tensor<2xi64>
  // CHECK: return %[[CST]]
}

// CHECK-LABEL: testFoldStridedSliceShapeDynamicOutput
func.func @testFoldStridedSliceShapeDynamicOutput(%arg0: tensor<?x1x2x?xf32>) -> (tensor<?xi32>) {

  %2 = mhlo.constant dense<0.000000e+00> : tensor<f32>

    -> (tensor<?x2xi64>, tensor<?x2xi64>, tensor<?x!tf_type.string>, tensor<?xi64>, tensor<2xi64>, tensor<2xi64>, tensor<32xf32>, tensor<32xf32>)

  // CHECK: tfrt.return {{.*}}, [[outputs]]#0
  func.return %outputs#0 : tensor<?x2xi64>
}

// CHECK-LABEL: func @no_native

// CHECK: return %[[GATHER]] : tensor<1x300x10xf32>

// -----

module {
  func.func @propagate_xlagather(%arg0: tensor<10x2xi32>) -> tensor<1x300x10xf32> {
    %cst = "tf.Const"() {value = dense<[1, 1, 300]> : tensor<3xi64>} : () -> tensor<3xi64>

  // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}>
  // CHECK: %[[R0:.*]] = "tf.Transpose"(%[[ARG0]], %[[CST]])
  // CHECK: %[[R1:.*]] = "tf.BiasAdd"(%[[R0]], %[[ARG1]]) <{data_format = "NHWC"}> {device = ""}
  // CHECK: %[[CST_0:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}>
  // CHECK: "tf.Transpose"(%[[R1]], %[[CST_0]])

  // CHECK-NEXT:  return [[FIND]] : tensor<*xi64>
  func.return %1 : tensor<*xi64>
}

// -----

// Test for case with int64 -> string hashtable.
func.func @hashtable_int64_to_string(%arg0: tensor<i64>) -> tensor<*x!tf_type.string> {
  %cst = arith.constant dense<"f"> : tensor<!tf_type.string>

// CHECK-LABEL: func @transposeConv2D
func.func @transposeConv2D(%input: tensor<1x32x32x3xf32>, %filter: tensor<1x1x3x8xf32>) -> tensor<1x7x7x8xf32> {

  // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}>
  // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]])

  // CHECK: %[[CONV2D:[0-9]*]] = "tf.Conv2D"(%[[ARG_TRANSPOSE]], %arg1)
  // CHECK-SAME: data_format = "NCHW"

  func.return %2, %3 : tensor<*xi64>, tensor<*xi64>
}

// -----

// CHECK-LABEL: @not_eliminate_multinomial
func.func @not_eliminate_multinomial(%0: tensor<*xf32>, %1: tensor<*xi32>) -> (tensor<*xi64>, tensor<*xi64>) {
  // CHECK-NEXT: tf.Multinomial
  // CHECK-SAME: seed = 0
  // CHECK-NEXT: tf.Multinomial
  // CHECK-SAME: seed = 1
  // CHECK-NEXT: tf.Multinomial