Speed up HMCEMAdaptiveStepSize... again.
Speed up HMCEMAdaptiveStepSize... again.

- Use tf.function to avoid growing the graph in graph mode. This removed the
  quadratic runtime, but increased the duration of each individual step, for
  unknown reason.
- Use TransformedVariable to take MCMC out of the GradientTape... somehow,
  despite all the stopped gradients it was still differentiating through the
  sampling loop. This made the graph test 4x faster.

PiperOrigin-RevId: 278975219

tensorflow/probability,tensorflow/probability