max value of a tensor in graph mode tensorflow

Question

I have this tf code in graph mode (it has a training function wrapped by @tf.function) where I need to get the max value of a tensor x with type

<class 'tensorflow.python.framework.ops.Tensor'> Tensor("x_3:0", shape=(100,), dtype=int64).

Then I need to use the max value of x as one of the arguments in the shape argument of tf.reshape(). If I print the output of tf.reduce_max(x) I get Tensor("Max:0", shape=(), dtype=int64), which is an invalid argument for tf.reshape(). I have tried tf.reduce_max(x).numpy() and it throws the error message 'Tensor' object has no attribute 'numpy'

So how do I get the max value of a tensor in graph mode in tf 2.6.0?

UPDATE This is my code with the necessary details (I hope) to see what is going on:

MyModel.py

class MyModel(tf.keras.Model):
    def __init__(self, ...,hidden_size,name='model',**kwargs):
        super(MyModel, self).__init__(name=name, **kwargs)
        self.hidden_size = hidden_size

    def call(self, inputs, training=True):
        x1, input, target, length, y = inputs
        batch_size = input.shape[0]
        print('check point 2', length, tf.reduce_max(length))
        padded_outputs = tf.reshape(tf.boolean_mask(outputs_dec,mask), shape=(batch_size,tf.reduce_max(length),self.hidden_size))
        print('check point 3',padded_outputs.shape)

    @tf.function
    def train(self, inputs, optimizer):
        with tf.GradientTape() as tape:
            costs = self.call(inputs)
        gradients = tape.gradient(self.loss, self.params)
        optimizer.apply_gradients(zip(gradients, self.params))

train_mymodel.py

tr_data = tf.data.Dataset.from_tensor_slices((x1_tr,                                             
                                              x2.input, 
                                              x2.target, 
                                              x2.length,
                                              y_tr))\                           
                                             .batch(args.batch_size)
while int(checkpoint.step) < args.epochs:
            for i, (x1_batch, input_batch, target_batch, length_batch, y_batch) in enumerate(tr_data):
                print('check point 1', length_batch)
                costs, obj_fn = mymodel.train((x1_batch, input_batch, target_batch, length_batch, y_batch),optimizer)

check point 1 tf.Tensor([300 300 ... 300 300], shape=(100,),type=int64)
check point 2 Tensor("x_3:0", shape=(100,), dtype=int64) Tensor("Max_1:0", shape=(), dtype=int64)
check point 3 (100, None, 500)

The shape of padded_outputs should be (100, 300, 500).

UPDATE2 The error happens when the graph is traced. If I hard code shape=(batch_size,300,self.hidden_size) and use tf.print(batch_size,tf.reduce_max(length),self.hidden_size) then the code runs without error messages and the output of tf.print() is (100,300,500). Is it any way to avoid such behavior?

Answer 1

It should work by simply passing the reduced tensor as an argument:

import tensorflow as tf
tf.random.set_seed(1)

@tf.function
def reshape_on_max_value():
  tensor1 = tf.random.uniform((5, 2), maxval=5, dtype=tf.int32)
  tensor2 = tf.random.uniform((4, 1), maxval=5, dtype=tf.int32)
  x = tf.reduce_max(tensor1)
  tf.print(type(tensor1), type(tensor2))
  tf.print(tf.reshape(tensor2, [x, 1, 1]).shape)

reshape_on_max_value()

<class 'tensorflow.python.framework.ops.Tensor'> <class 'tensorflow.python.framework.ops.Tensor'>
TensorShape([4, 1, 1])