Excessive retracing in autobatched joint distributions

Calling .sample or .log_prob repeatedly on an autobatched joint distribution (as you would in, say, an eager-mode optimization loop) prints the warning

WARNING:tensorflow:5 out of the last 5 calls to <function pfor.<locals>.f at 0x7f16b7927e18> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args and https://www.tensorflow.org/api_docs/python/tf/function for more details.

Reproduction:
https://colab.research.google.com/drive/15RRL7pjgORFzcjWH2i74BQte61EUtBGc?usp=sharing

Autobatched JDs should probably cache the pfor'd tf.functions to avoid this.

Start with the linked Colab reproduction and follow the autobatched joint distribution paths used by repeated sample and log_prob calls. Inspect where the pfor'd tf.functions are created, then verify that repeated calls no longer emit excessive retracing warnings. Done means the reproduction runs without the warning while preserving the distribution behavior.