from tensorflow.quantization import fake_quant_with_min_max_vars
from tensorflow.python.framework import graph_util
import math
import numpy as np
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
from tensorflow.python.keras.layers import dense_attention

np.set_printoptions(threshold=np.inf)
np.random.seed(0)
input_shape = [(1, 2, 2, 5), (1, 2, 2, 3), (1, 1, 3, 5), (2, 3, 2, 3), (1, 13, 13, 6), (2, 2, 2, 14)]
axis_list = (0, 1, 2, 3, 3, 1)
idx = 5
axis = axis_list[idx]
input = np.random.randint(0, 16384, input_shape[idx]).astype(dtype=np.float32)
# input = np.random.rand(*input_shape[0]).astype(dtype=np.float32)
# input = np.array([10, 20, 30, 40, 90, 20, 11, 15, 1, 50, 30, 45, 25, 50, 30]).reshape(input_shape[2]).astype(dtype=np.float32)
# input = np.array([100, 2,  300, 4,  50, 6,  11, 12, 13, 34, 35, 36,  9,  6, 17, 10, 20, 30,
#                   10,  20, 30,  40, 5,  60, 7,  80, 90, 10, 11, 120, 18, 5, 16, 9,  22, 23]).reshape(input_shape[3]).astype(dtype=np.float32)
print("np input: ", input.shape, input.flatten())

output = tf.argmax(input, axis, output_type=tf.int64, name=None)
# output = tf.reduce_max(input, 1, name=None, keepdims=True)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    output_gnd = output.eval(session=sess)
print("tf output shape: ", output_gnd.shape, "output data: ", output_gnd)