FunMC: Add a name scope for the hmc_with_state_grads.

SiegeLordEx · tensorflower-gardener · commit 827ea3db6b4b · 2021-06-03T14:26:56.000-07:00
PiperOrigin-RevId: 377372672
diff --git a/spinoffs/fun_mc/fun_mc/prefab.py b/spinoffs/fun_mc/fun_mc/prefab.py
@@ -662,6 +662,7 @@ class HamiltonianMonteCarloWithStateGradsExtra(NamedTuple):
   proposed_state: 'fun_mc.State'
 
 
+@util.named_call
 def hamiltonian_monte_carlo_with_state_grads_step(
     hmc_state: 'fun_mc.HamiltonianMonteCarloState',
     trajectory_length: 'fun_mc.FloatTensor',
