Memory leak for TFP distribution

Hi,

I'm facing a memory leak issue with TFP distributions. Specifically, I have a custom function that samples from a Categorical distribution and returns the output:

def compute_categorical_distribution(classes_pred,dtype=np.int8,validate_args=True):
  # Define the distribution
  distr = tfp.distributions.Categorical(probs=classes_pred, validate_args=validate_args)
  # Sample from the distribution
  output = distr.sample().numpy().astype(dtype)
del distr
K.clear_session()
gc.collect()
return output

I was expecting that, when I call this function, I would get the output array back and the memory allocation of the distribution would be freed up. However, this does not seem to be the case, as for every consecutive call of this function the memory usage keeps increasing, indicating that a new tfp distribution is created each time and remains in the memory. Is there a way to clear up the memory allocation of the distribution? As you can see above, I tried del distr, tensorflow.keras.backend.clear_session() and gc.collect() but none of these seems to make any difference. I am using tfp version 0.19.0.

Thanks!
Fotis

Run the provided compute_categorical_distribution function repeatedly with TFP 0.19.0 while measuring process memory; begin with the Categorical construction and the del, clear_session, and gc.collect sequence shown. Done means determining why memory grows and verifying whether repeated calls release the distribution allocations.