NUTS and array of log-likelihoods

Can I use a theano operation that evaluates to a vector as the log-likelihood for NUTS? In this case, would theano try to optimize each element of the array of log-likelihood independently?

I want to know if NUTS can do this for the following reason. The calculation of my log-likelihood requires iteration over the whole data set. The current iteration depends on the previous iteration. If I try to calculate 1000 log-likelihood for 1000 sets of parameters parallely, I might be able to use the GPU to increase throughput.

Tensorflow’s Hamiltonian Monte Carlo can do that. But Tensorflow doesn’t have NUTS.
https://www.tensorflow.org/probability/api_docs/python/tfp/mcmc/HamiltonianMonteCarlo

I guess even if NUTS does not support this feature, I can work around it in principle.

The overall model contains 1000 replicas that are identical except having different values for the parameters.

Since all replicas are statistically (but not computationally) independent, the overall log-likelihood is sum of log-likelihoods of all replicas.

Well … now NUTS knows the sum of the log-likelihoods of all replicas but not the log-likelihood of each replica. Would NUTS become slower when it doesn’t know the log-likelihood of each replica?