An elegant replacement for numpy.random.choice in Tensorflow in this context

I have been programming a machine learning pipeline in Tensorflow 2.

Below is a code snippet implementing the sampling from a replay buffer in a policy gradient algorithm. A code similar to mine can be seen here. I feel the rest of the code would be too long and unnecessary for this question.

    def sample_buffer(self, batch_size):
        max_mem = tf.math.minimum(self.mem_cntr, self.mem_size)

        batch = np.random.choice(max_mem, batch_size)
        batch = tf.convert_to_tensor(batch, dtype=tf.int32)

        states = tf.gather(self.state_memory, batch)
        states_ = tf.gather(self.new_state_memory, batch)
        actions = tf.gather(self.action_memory, batch)
        rewards = tf.gather(self.reward_memory, batch)

        return states, actions, rewards, states_

This piece of code causes this error during execution - "NotImplementedError: Cannot convert a symbolic tf.Tensor (cond/Minimum:0) to a numpy array. This error may indicate that you're trying to pass a Tensor to a NumPy call, which is not supported."

I initialised the variables mem_cntr and mem_size as Pythonic integers, and originally the tf.math.minimum was simply a min, which gave an error ("OperatorNotAllowedInGraphError: Using a symbolic tf.Tensor as a Python bool is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.") Hence the use of tf.math.minimum even though I never initialised these two variables as Tensors.

I tried this code suggested in another Stackoverflow post. This code works in the Python interpreter in my terminal.

def select_indices_uniform(dims, num_indices, unique=False):
    # Create a tensor of probabilities
    size = tf.reduce_prod(tf.constant(dims, dtype=tf.int32)).numpy()

    # Use the uniform_candidate_sampler function to sample indices
    samples = tf.random.log_uniform_candidate_sampler(
            true_classes=[[0]], num_true=1, num_sampled=num_indices, unique=unique, range_max=size)

    # Extract the data and return it
    indices = tf.unravel_index(samples[0], dims)
    return indices

But the problem here is that eager execution does not work with my pipeline (so I can't use .numpy()), and using the Tensorflow v1.x compatible eval to try to evaluate the tensor to a numpy array fails due to a Graph execution error.

I can not think of another way to draw a fixed number of samples from these 4 tensors without this or numpy.random.choice. tf.boolean_mask and tf.random.categorical would be good ways to draw a random number of samples.

Is there any workaround?

Please help!