I’m sure that any robust PR implementing new features would be very welcome.
I just wanted to check that you were aware of PyMC4? One of the aims seems to be better GPU usage via tensorflow. I imagine that a biproduct of this would be better multi-node usage as well. I mention it in case it is helpful .