LogicalResult MatchReduceToArgMinMaxType2(mhlo::ReduceOp reduce_op,
                                          bool is_argmax);

// Base class for converting mhlo::ReduceOp to TF/TFL ArgMax/ArgMin ops.
template <typename Reduce, typename ArgReduce, typename BooleanReduce,
          bool is_argmax>
class ConvertReduceOpToArgMinMax : public OpConversionPattern<mhlo::ReduceOp> {
 public:

#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {
// Pattern matches the following reduction function for ArgMax/ArgMin coming
// from PyTorch
// %0 = compare{GT}(%lhs_value, %rhs_value)
// %1 = select(%0, %lhs_value, %rhs_value)
// %2 = compare{EQ}(%lhs_value, %rhs_value)
// %3 = minimum(%lhs_index, %rhs_index)

      loss_value = tf.reduce_mean(
          tf.nn.softmax_cross_entropy_with_logits(labels, logits))

    grads = tape.gradient(loss_value, model.trainable_variables)
    correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(labels, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return accuracy, loss_value

  @tf.function

Number of clusters: 1

unclustered size 446
 AssignAddVariableOp 2
 Const 2
 DivNoNan 1
 Identity 1
 NoOp 1
 ReadVariableOp 2
 VarHandleOp 435
 _Retval 2
cluster 0 size 2178
 Add 17
 AddN 72
 ArgMax 1
 AssignAddVariableOp 1
 AssignSubVariableOp 106
 BiasAdd 1
 BiasAddGrad 1
 Cast 3
 Const 357
 Conv2D 53
 Conv2DBackpropFilter 53
 Conv2DBackpropInput 52
 DivNoNan 1
 Equal 1
 FusedBatchNorm 53

unclustered size 606
 AssignAddVariableOp 1
 Const 108
 DivNoNan 1
 Identity 2
 Merge 53
 PlaceholderWithDefault 1
 ReadVariableOp 2
 Switch 1
 _Arg 435
 _Retval 2
cluster 0 size 1910
 Add 16
 AddN 71
 ArgMax 1
 AssignAddVariableOp 1
 BiasAdd 1
 BiasAddGrad 1
 Cast 115
 Const 407
 Conv2D 53
 Conv2DBackpropFilter 53
 Conv2DBackpropInput 52
 Equal 1
 FusedBatchNormGradV2 53
 FusedBatchNormV2 53
 MatMul 3

    ],
    data = [
        "//tensorflow/lite/tools/optimize:testdata/add_with_const_input.bin",
        "//tensorflow/lite/tools/optimize:testdata/argmax.bin",
        "//tensorflow/lite/tools/optimize:testdata/broadcast_to.bin",
        "//tensorflow/lite/tools/optimize:testdata/concat.bin",
        "//tensorflow/lite/tools/optimize:testdata/fc.bin",

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

	{"hw.usermem", []_C_int{6, 20}},
	{"hw.uuid", []_C_int{6, 18}},
	{"hw.vendor", []_C_int{6, 14}},
	{"hw.version", []_C_int{6, 16}},
	{"kern.allowdt", []_C_int{1, 65}},
	{"kern.allowkmem", []_C_int{1, 52}},
	{"kern.argmax", []_C_int{1, 8}},
	{"kern.audio", []_C_int{1, 84}},
	{"kern.boottime", []_C_int{1, 21}},
	{"kern.bufcachepercent", []_C_int{1, 72}},
	{"kern.ccpu", []_C_int{1, 45}},
	{"kern.clockrate", []_C_int{1, 12}},

	{"hw.usermem", []_C_int{6, 20}},
	{"hw.uuid", []_C_int{6, 18}},
	{"hw.vendor", []_C_int{6, 14}},
	{"hw.version", []_C_int{6, 16}},
	{"kern.allowdt", []_C_int{1, 65}},
	{"kern.allowkmem", []_C_int{1, 52}},
	{"kern.argmax", []_C_int{1, 8}},
	{"kern.audio", []_C_int{1, 84}},
	{"kern.boottime", []_C_int{1, 21}},
	{"kern.bufcachepercent", []_C_int{1, 72}},
	{"kern.ccpu", []_C_int{1, 45}},
	{"kern.clockrate", []_C_int{1, 12}},

	{"hw.usermem", []_C_int{6, 20}},
	{"hw.uuid", []_C_int{6, 18}},
	{"hw.vendor", []_C_int{6, 14}},
	{"hw.version", []_C_int{6, 16}},
	{"kern.allowdt", []_C_int{1, 65}},
	{"kern.allowkmem", []_C_int{1, 52}},
	{"kern.argmax", []_C_int{1, 8}},
	{"kern.audio", []_C_int{1, 84}},
	{"kern.boottime", []_C_int{1, 21}},
	{"kern.bufcachepercent", []_C_int{1, 72}},
	{"kern.ccpu", []_C_int{1, 45}},
	{"kern.clockrate", []_C_int{1, 12}},