Add time values as sampler stats for NUTS

[aseyboldt]

This PR adds three more sampler stats for NUTS and HMC: "process_time_diff_ns" and "perf_counter_diff_ns" and "perf_counter_ns".
During debugging it can be useful to see how long it took to compute each sample, and for issues involving blas/openmp it can also be useful to see the difference between process time and wall time.
The names are taken from the python time module: https://docs.python.org/3/library/time.html#time.perf_counter

This can be used like this:

with pm.Model() as model:
    pm.Normal("a", shape=1000)
    
    tr = pm.sample(return_inferencedata=True)

stats = tr.sample_stats
for chain in stats.chain:
    sns.distplot(np.log(stats.process_time_diff_ns.sel(chain=chain)), label=chain.values)
plt.legend();

[AlexAndorra]

Neat, thanks @aseyboldt ! Left some comments below

[aseyboldt]

I forgot to check when the time functions were added to the stdlib. The nanosecond versions are new in 3.7. We still want to support 3.6 for some time?
I guess we can just switch to the floating point versions then. The resolution should still be good enough I think.

[AlexAndorra]

Just a couple questions below before approving

[AlexAndorra]

Just a suggestion: why not call this perf_time_diff? I find it more explicit and it matches process_time_diff

[aseyboldt]

I see why perf_time_diff might be nicer, I just followed the naming of the function in time, so that it's easier to see what clock is used exactly. But counter is a bit confusing...

[AlexAndorra]

Same comment on the name, and out of curiosity: why did switch from perf_end to perf_start?

[aseyboldt]

Again, I think the choice it kind of arbitrary. We can reconstruct the other one since we have the difference.
I just thought it might be a bit more intuitive to have the start of the draw as an absolute value.

[junpenglao]

Is it something we should considered push to top level api so that it is available for all step_methods?

https://github.com/pymc-devs/pymc3/blob/747db63948f8115e30d676089b77116791a028fa/pymc3/step_methods/arraystep.py#L145-L157

FWIW, I think it is fine for it to be a HMC only method, and change later if there are feature request - just want to bring up this point.

[twiecki]

Also needs note in release notes.

[aseyboldt]

@junpenglao That would be useful, but we don't have a way to easily add sampler stats to all samplers at once, since they have to be declared in the step method itself. If users implement their own step methods (does anyone?) that would be a breaking change. At least unless we change some code in the trace backends to allow for missing or undeclared stats.

@twiecki done

[AlexAndorra]

Sorry, just one last nitpick 😜

[twiecki]

I'm also fine with dropping Python 3.6, I know it's just a small thing here with an easy work-around but I think being progressive here is a good thing.

[ColCarroll]

Good news about 1 week ago, then (via https://numpy.org/neps/nep-0029-deprecation_policy.html):

Date	Python	NumPy
Jan 07, 2020	3.6+	1.15+
Jun 23, 2020	3.7+	1.15+
Jul 23, 2020	3.7+	1.16+
Jan 13, 2021	3.7+	1.17+
Jul 26, 2021	3.7+	1.18+
Dec 26, 2021	3.8+	1.18+
Apr 14, 2023	3.9+	1.18+

[junpenglao]

@aseyboldt I see.

+1 to dropping py3.6

[twiecki]

Great, let's drop it then!

…

On Wed, Jul 1, 2020 at 4:34 PM Junpeng Lao ***@***.***> wrote: @aseyboldt <https://github.com/aseyboldt> I see. +1 to dropping py3.6 — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#3986 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAFETGAYNIAJT64T6MCRK2TRZNCPDANCNFSM4ONFUG3Q> .

[aseyboldt]

Dropping 3.6 is fine, but maybe the float version is better after all. :-)
It has unit 'second' which is easier to work with, and we don't need good resolution in the ns range, we just aren't that fast.

[AlexAndorra]

All good now, thanks @aseyboldt ! I agree with you on the seconds vs. nanoseconds stuff

Problem is fixed now and this is blocking merge