text: Title: `nuts.get_extra_fields()["num_steps"]=0` after warmup
  Body: I have the following piece of code:
  ``` python
  nuts = MCMC(
      NUTS(model_logreg),
      num_warmup=2**13,
      num_samples=2**10,
      num_chains=2**5,
      chain_method="vectorized",
  )
  nuts.warmup(jr.key(2), x_train, labels_train, extra_fields=("num_steps",))
  warmup_steps = nuts.get_extra_fields()["num_steps"]
  print(f"num warmup steps: {warmup_steps}")
  ```
  which returns
  ```
  warmup: 100%|██████████| 8192/8192 [00:41<00:00, 199.13it/s]
  num warmup steps: [0 0 0 ... 0 0 0]
  ```
  
  If I do `nuts.run(jr.key(2), x_train, labels_train, extra_fields=("num_steps",)` it works just fine and reports a non-zero number of steps (although I suspect it doesn't count the warmup steps). Also the sampling itself works as intended and results in the correct distribution, so the problem probably isn't in my code. And the warmup does indeed work, because if I set `num_warmup=0`, then the output becomes biased towards the initial value.
  
  This is quite bad because it makes it seem that NUTS can achieve good results with a very small number of gradient evaluations, giving it an unfair advantage over other samplers.
  
  Also I saw this issue mentioned in the following thread, but it apparently hasn't been addressed yet:
  https://forum.pyro.ai/t/how-to-calculate-effective-sample-size-per-gradient-evaluation/5398/7