Use float time counters for nuts stats

pymc-devs · Jul 1, 2020 · b60d81c · b60d81c
1 parent 334dfe1
commit b60d81c
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 19 deletions.
diff --git a/pymc3/step_methods/hmc/base_hmc.py b/pymc3/step_methods/hmc/base_hmc.py
@@ -133,8 +133,8 @@ def _hamiltonian_step(self, start, p0, step_size):
 
     def astep(self, q0):
         """Perform a single HMC iteration."""
-        perf_start = time.perf_counter_ns()
-        process_start = time.process_time_ns()
+        perf_start = time.perf_counter()
+        process_start = time.process_time()
 
         p0 = self.potential.random()
         start = self.integrator.compute_state(q0, p0)
@@ -170,8 +170,8 @@ def astep(self, q0):
 
         hmc_step = self._hamiltonian_step(start, p0, step_size)
 
-        perf_end = time.perf_counter_ns()
-        process_end = time.process_time_ns()
+        perf_end = time.perf_counter()
+        process_end = time.process_time()
 
         self.step_adapt.update(hmc_step.accept_stat, adapt_step)
         self.potential.update(hmc_step.end.q, hmc_step.end.q_grad, self.tune)
@@ -201,9 +201,9 @@ def astep(self, q0):
         stats = {
             "tune": self.tune,
             "diverging": bool(hmc_step.divergence_info),
-            "perf_counter_diff_ns": perf_end - perf_start,
-            "process_time_diff_ns": process_end - process_start,
-            "perf_counter_ns": perf_end,
+            "perf_counter_diff": perf_end - perf_start,
+            "process_time_diff": process_end - process_start,
+            "perf_counter_start": perf_start,
         }
 
         stats.update(hmc_step.stats)

diff --git a/pymc3/step_methods/hmc/hmc.py b/pymc3/step_methods/hmc/hmc.py
@@ -48,9 +48,9 @@ class HamiltonianMC(BaseHMC):
         'path_length': np.float64,
         'accepted': np.bool,
         'model_logp': np.float64,
-        'process_time_diff_ns': np.int64,
-        'perf_counter_diff_ns': np.int64,
-        'perf_counter_ns': np.int64,
+        'process_time_diff': np.float64,
+        'perf_counter_diff': np.float64,
+        'perf_counter_start': np.float64,
     }]
 
     def __init__(self, vars=None, path_length=2., max_steps=1024, **kwargs):

diff --git a/pymc3/step_methods/hmc/nuts.py b/pymc3/step_methods/hmc/nuts.py
@@ -72,13 +72,13 @@ class NUTS(BaseHMC):
       samples, the step size is set to this value. This should converge
       during tuning.
     - `model_logp`: The model log-likelihood for this sample.
-    - `process_time_diff_ns`: The time it took to draw the sample, as defined
-      by the python standard library `time.process_time_ns`. This counts all
+    - `process_time_diff`: The time it took to draw the sample, as defined
+      by the python standard library `time.process_time`. This counts all
       the CPU time, including worker processes in BLAS and OpenMP.
-    - `perf_counter_diff_ns`: The time it took to draw the sample, as defined
-      by the python standard library `time.perf_counter_ns` (wall time).
-    - `perf_counter_ns`: The value of the `time.perf_counter_ns` after drawing
-      the sample.
+    - `perf_counter_diff`: The time it took to draw the sample, as defined
+      by the python standard library `time.perf_counter` (wall time).
+    - `perf_counter_start`: The value of `time.perf_counter` at the beginning
+      of the computation of the draw.
 
     References
     ----------
@@ -103,9 +103,9 @@ class NUTS(BaseHMC):
             "energy": np.float64,
             "max_energy_error": np.float64,
             "model_logp": np.float64,
-            "process_time_diff_ns": np.int64,
-            "perf_counter_diff_ns": np.int64,
-            "perf_counter_ns": np.int64,
+            "process_time_diff": np.float64,
+            "perf_counter_diff": np.float64,
+            "perf_counter_start": np.float64,
         }
     ]