Optimise function for many pseudodata realisations in TensorFlow 2

My end goal is to simulate likelihood ratio test statistics, however the core problem I am having is that I do not understand how to get TensorFlow 2 to perform many optimisations for different data inputs. Here is my attempt, hopefully it gives you the idea of what I am trying:

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd
import numpy as np

# Bunch of independent Poisson distributions that we want to combine
poises0 = [tfp.distributions.Poisson(rate = 10) for i in range(5)]

# Construct joint distributions
joint0 = tfd.JointDistributionSequential(poises0)

# Generate samples
N = int(1e3)
samples0 = joint0.sample(N)

# Now we need the same distributions but with floating parameters,
# and need to define the function to be minimised
mus = [tf.Variable(np.random.randn(), name='mu{0}'.format(i)) for i in range(5)]

#@tf.function
def loss():
    poises_free = [tfp.distributions.Poisson(rate = mus[i]) for i in range(5)]
    joint_free = tfd.JointDistributionSequential(poises_free)
    # Construct (half of) test statistic
    return -2*(joint_free.log_prob(samples0))

# Minimise (for all samples? Apparently not?)
opt = tf.optimizers.SGD(0.1).minimize(loss,var_list=mus)

print(mus)
print(loss())
print(opt)
quit()

Output:

[<tf.Variable 'mu0:0' shape=() dtype=float32, numpy=53387.016>, <tf.Variable 'mu1:0' shape=() dtype=float32, numpy=2540.568>, <tf.Variable 'mu2:0' shape=() dtype=float32, numpy=-5136.6226>, <tf.Variable 'mu3:0' shape=() dtype=float32, numpy=-3714.5227>, <tf.Variable 'mu4:0' shape=() dtype=float32, numpy=1062.9396>]
tf.Tensor(
[nan nan nan nan ... nan nan nan], shape=(1000,), dtype=float32)
<tf.Variable 'UnreadVariable' shape=() dtype=int64, numpy=1>

In the end I want to compute the test statistic

q = -2*joint0.log_prob(samples0) - loss()

and show that it has a chi-squared distribution with 5 degrees of freedom.

I am new to TensorFlow so perhaps I am doing this entirely wrong, but I hope you get the idea of what I want.

Edit:

So I played around a bit more, and I suppose that TensorFlow simply doesn't perform optimisations over the input tensors in parallel like I assumed. Or perhaps it can, but I need to set things up differently, i.e. perhaps give it a tensor of input parameters and a gigantic joint loss function for all the minimisations at once?

I also tried doing things with a simple loop just to see what happens. As predicted it is pathetically slow, but I also don't even get the right answer:

poises0 = [tfp.distributions.Poisson(rate = 10) for i in range(5)]
joint0 = tfd.JointDistributionSequential(poises0)

N = int(5e2)
samples0 = joint0.sample(N)

mus = [tf.Variable(10., name='mu{0}'.format(i)) for i in range(5)]

#@tf.function
def loss(xi):
    def loss_inner():
        poises_free = [tfp.distributions.Poisson(rate = mus[i]) for i in range(5)]
        joint_free = tfd.JointDistributionSequential(poises_free)
        # Construct (half of) test statistic
        return -2*(joint_free.log_prob(xi))
    return loss_inner

# Minimise
# I think I have to loop over the samples... bit lame. Can perhaps parallelise though.
q = []
for i in range(N):
   xi = [x[i] for x in samples0]
   opt = tf.optimizers.SGD(0.1).minimize(loss=loss(xi),var_list=mus)
   q += [-2*joint0.log_prob(xi) - loss(xi)()]

fig = plt.figure()
ax = fig.add_subplot(111)
sns.distplot(q, kde=False, ax=ax, norm_hist=True)
qx = np.linspace(np.min(q),np.max(q),1000)
qy = np.exp(tfd.Chi2(df=5).log_prob(qx))
sns.lineplot(qx,qy)
plt.show()

The output is not a chi-squared distribution with DOF=5. Indeed the test statistic often has negative values, which means that the optimised result is often a worse fit than the null hypothesis, which should be impossible.

Edit 2:

Here is an attempt at the "monster" solution where I minimise a giant network of different input variables for each pseudodata realisation all at once. This feels more like something that TensorFlow might be good at doing, though I feel like I will run out of RAM once I go to large sets of pseudodata. Still, I can probably loop over batches of pseudodata.

poises0 = [tfp.distributions.Poisson(rate = 10) for i in range(5)]
joint0 = tfd.JointDistributionSequential(poises0)

N = int(5e3)
samples0 = joint0.sample(N)

mus = [tf.Variable(10*np.ones(N, dtype='float32'), name='mu{0}'.format(i)) for i in range(5)]

poises_free = [tfp.distributions.Poisson(rate = mus[i]) for i in range(5)]
joint_free = tfd.JointDistributionSequential(poises_free)
qM = -2*(joint_free.log_prob(samples0))

@tf.function
def loss():
    return tf.math.reduce_sum(qM,axis=0)

# Minimise
opt = tf.optimizers.SGD(0.1).minimize(loss,var_list=mus)
print("parameters:", mus)
print("loss:", loss())
q0 =-2*joint0.log_prob(samples0)
print("q0:", q0)
print("qM:", qM)
q = q0 - qM

fig = plt.figure()
ax = fig.add_subplot(111)
sns.distplot(q, kde=False, ax=ax, norm_hist=True)
qx = np.linspace(np.min(q),np.max(q),1000)
qy = np.exp(tfd.Chi2(df=5).log_prob(qx))
sns.lineplot(qx,qy)
plt.show()

Unfortunately I now get the error:

Traceback (most recent call last):
  File "testing3.py", line 35, in <module>
    opt = tf.optimizers.SGD(0.1).minimize(loss,var_list=mus)   
  File "/home/farmer/anaconda3/envs/general/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py", line 298, in minimize
    return self.apply_gradients(grads_and_vars, name=name)
  File "/home/farmer/anaconda3/envs/general/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py", line 396, in apply_gradients
    grads_and_vars = _filter_grads(grads_and_vars)
  File "/home/farmer/anaconda3/envs/general/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py", line 924, in _filter_grads
    ([v.name for _, v in grads_and_vars],))
ValueError: No gradients provided for any variable: ['mu0:0', 'mu1:0', 'mu2:0', 'mu3:0', 'mu4:0'].

which I suppose is a basic sort of error. I think I just don't understand how TensorFlow keeps track of the derivatives it needs to compute. It seems like things work if I define variables inside the loss function rather than outside, but I need them outside in order to access their values later. So I guess I don't understand something here.

from Optimise function for many pseudodata realisations in TensorFlow 2