// CHECK-NEXT:  %[[argmax:.*]] = "tfl.arg_max"(%[[softmax]], %[[cst]]) : (tensor<128x16xf32>, tensor<i32>) -> tensor<128xi32>
  %3 = "tfl.arg_max"(%2, %cst) : (tensor<128x16xf32>, tensor<i32>) -> tensor<128xi32>
// CHECK-NEXT:  return %[[softmax]], %[[argmax]] : tensor<128x16xf32>, tensor<128xi32>
  func.return %2, %3 : tensor<128x16xf32>, tensor<128xi32>
}

// CHECK-LABEL: PruneUnusedLstm

  %1 = "tfl.pseudo_qconst"() {qtype = tensor<128x!quant.uniform<i32:f32, 0.7>>, value = dense<0> : tensor<128xi32>} : () -> tensor<128x!quant.uniform<i32:f32, 0.7>>

// CHECK:         }
}

// -----

module {

        }) : () -> tensor<f32>
        "tf_device.cluster"() ({
          %group = "tf.Const"() <{value = dense<[[0, 1]]> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
          %arg1_reduced = "tf.XlaAllReduce"(%arg1_id, %group) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<f32>, tensor<1x2xi32>) -> tensor<f32>
          "tf.OpA"(%arg1_reduced) : (tensor<f32>) -> ()
          tf_device.return
        }) : () -> ()

  func.func private @uniform_quantize_0(%arg0: tensor<1x2xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x2xi32> {
    %0 = stablehlo.convert %arg0 : (tensor<1x2xf32>) -> tensor<1x2xi32>
    return %0 : tensor<1x2xi32>
  }
// CHECK: @uniform_quantize_0
  func.func private @uniform_quantize_1(%arg0: tensor<1x3xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x3xi8> {

// CHECK-NO-UNPACK-SAME: (%[[ARG_0:.+]]: tensor<1x1024xf32>) -> tensor<1x3xf32>
// CHECK-NO-UNPACK-DAG: %[[CONST:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1024x3xi8>}> : () -> tensor<1024x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>>
// CHECK-NO-UNPACK: %[[QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x1024xf32>) -> tensor<1x1024x!quant.uniform<i8:f32, {{.*}}>>

    }
  }
)mlir";

constexpr absl::string_view kModuleDynamic = R"mlir(
  module {
    func.func @main(%arg0: tensor<?x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<?x3xf32> attributes {_from_xla_call_module} {
      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [] : (tensor<?x1024xf32>, tensor<1024x3xf32>) -> tensor<?x3xf32>
      return %0 : tensor<?x3xf32>
    }
  }
)mlir";

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

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

  %0 = "tfl.pseudo_qconst"() {qtype = tensor<128x!quant.uniform<i32:f32, 0.7>>, value = dense<0> : tensor<128xi32>} : () -> tensor<128x!quant.uniform<i32:f32, 0.7>>

func.func @reshape_vector_shape(tensor<4x4x4xf32>) -> tensor<16x4xf32> {
^bb0(%arg0: tensor<4x4x4xf32>) :
  %shape0 = arith.constant dense<[[16, 4]]> : tensor<1x2xi32>
  // expected-error @+1 {{'tfl.reshape' op requires 'shape' to be rank 1, but got 2}}
  %1 = "tfl.reshape"(%arg0, %shape0) : (tensor<4x4x4xf32>, tensor<1x2xi32>) -> tensor<16x4xf32>
  func.return %1 : tensor<16x4xf32>
}

// -----