Why does tf.random.categorical gives a None dimension when num_samples is batch-dependent?

I am making a model that involves random numbers generations (tf.random.categorical / tf.random.gamma).

The num_samples parameter I use have to be batch-dependent (they are the maximum values of variables which first dimensions is the batch_size).

The issue is that the result of those tf.random functions is always a matrix which second dimension is None. Which is annoying because I need to pass this dimension size into functions that do not accept None values).

In this code below, I made a simple model that reproduces this issue.

import tensorflow as tf
import numpy as np
batch_size = 5
input_length = 111

class MyLayer(tf.keras.layers.Layer):
    def __init__(self,input_length,**kwargs):
        self.state_size = [tf.TensorShape([input_length])]
        super(MyLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.built = True

    def call(self, inputs, states):
        Pro = tf.math.log(1/50*tf.ones(50))
        size_batch = len(inputs)
        
        #Works well when the sample number is a constant
        num_samples = tf.constant(10)
        result_matrix = tf.cast(tf.random.categorical(tf.repeat([Pro],size_batch,0),num_samples),tf.float32)
        print('result_matrix if num_samples is constant: ',result_matrix)
        
        #Shape 1 is None when the sample number is batch-dependent
        num_samples = int(tf.reduce_max(inputs))
        result_matrix = tf.cast(tf.random.categorical(tf.repeat([Pro],size_batch,0),num_samples),tf.float32)
        print('result_matrix if num_samples is batch-dependent: ',result_matrix)
        
        states = inputs
        output = inputs[:,0]
        return output, states

cell = MyLayer(input_length)
layer = tf.keras.layers.RNN(cell)

inp1 = tf.keras.Input(shape=(None,input_length))
rnn = tf.keras.layers.RNN(cell, return_state=True)
output = rnn(inp1) 
model = tf.keras.models.Model(inp1, output)

x = tf.ones((100,10,111))
y = tf.ones((100))
model.compile(optimizer="adam", loss="mse", metrics=["accuracy"])
model.fit(x,y,batch_size = batch_size)

Output:

2022-01-09 16:39:00.216815: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cudart64_110.dll'; dlerror: cudart64_110.dll not found
2022-01-09 16:39:00.231292: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
result_matrix if num_samples is constant:  Tensor("rnn_1/my_layer/Cast:0", shape=(None, 10), dtype=float32)
result_matrix if num_samples is batch-dependent:  Tensor("rnn_1/my_layer/Cast_2:0", shape=(None, None), dtype=float32)
result_matrix if num_samples is constant:  Tensor("rnn_1/while/my_layer/Cast:0", shape=(None, 10), dtype=float32)
result_matrix if num_samples is batch-dependent:  Tensor("rnn_1/while/my_layer/Cast_2:0", shape=(None, None), dtype=float32)
2022-01-09 16:39:03.182967: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cudart64_110.dll'; dlerror: cudart64_110.dll not found
2022-01-09 16:39:03.188726: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cublas64_11.dll'; dlerror: cublas64_11.dll not found
2022-01-09 16:39:03.194037: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cublasLt64_11.dll'; dlerror: cublasLt64_11.dll not found
2022-01-09 16:39:03.198963: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cufft64_10.dll'; dlerror: cufft64_10.dll not found
2022-01-09 16:39:03.204394: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'curand64_10.dll'; dlerror: curand64_10.dll not found
2022-01-09 16:39:03.209105: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cusolver64_11.dll'; dlerror: cusolver64_11.dll not found
2022-01-09 16:39:03.214256: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cusparse64_11.dll'; dlerror: cusparse64_11.dll not found
2022-01-09 16:39:03.219299: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cudnn64_8.dll'; dlerror: cudnn64_8.dll not found
2022-01-09 16:39:03.223291: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1850] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
2022-01-09 16:39:03.232555: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX AVX2
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
result_matrix if num_samples is constant:  Tensor("model/rnn_1/my_layer/Cast:0", shape=(5, 10), dtype=float32)
result_matrix if num_samples is batch-dependent:  Tensor("model/rnn_1/my_layer/Cast_2:0", shape=(5, None), dtype=float32)
result_matrix if num_samples is constant:  Tensor("model/rnn_1/while/my_layer/Cast:0", shape=(5, 10), dtype=float32)
result_matrix if num_samples is batch-dependent:  Tensor("model/rnn_1/while/my_layer/Cast_2:0", shape=(5, None), dtype=float32)
result_matrix if num_samples is constant:  Tensor("model/rnn_1/my_layer/Cast:0", shape=(5, 10), dtype=float32)
result_matrix if num_samples is batch-dependent:  Tensor("model/rnn_1/my_layer/Cast_2:0", shape=(5, None), dtype=float32)
result_matrix if num_samples is constant:  Tensor("model/rnn_1/while/my_layer/Cast:0", shape=(5, 10), dtype=float32)
result_matrix if num_samples is batch-dependent:  Tensor("model/rnn_1/while/my_layer/Cast_2:0", shape=(5, None), dtype=float32)
 1/20 [>.............................] - ETA: 7s - loss: 0.0000e+00 - rnn_1_loss: 0.0000e+00 - rnn_1_1_loss: 0.0000e+00 20/20 [==============================] - 0s 821us/step - loss: 0.0000e+00 - rnn_1_loss: 0.0000e+00 - rnn_1_1_loss: 0.0000e+00 - rnn_1_accuracy: 1.0000 - rnn_1_1_accuracy: 0.0000e+00
>>>

To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags. Does anyone know why this is happening?

Does anyone know what I could do so that I don't have this None dimension issue?