Add a module to apply updates every k steps (and accumulate them othe…

[perolat]

…rwise)

[mtthss]

Looks great!

[tomhennigan]

The name reset doesn't actually reflect what's in the variable and I think the guard is wrong (afaik reset == 0 should actually be reset != 0).

I'd suggest:

c = state.count % k
acc = c != 0
grad_acc = tree_multimap(lambda g, ga: acc * ga + g, updates, state.grad_acc)
emit = c == (k - 1)
updates = tree_multimap(lambda ga: emit * ga, grad_acc)

It would also be great if there was a test for this to avoid regressions..

[perolat]

Thanks Tom for spotting that problem. I'd be happy to do some tests but where is optix tested?

[mtthss]

https://github.com/google/jax/blob/master/tests/optix_test.py

[perolat]

Tests added!

The tolerance is quite high absolute tolerance is 1e-6 and the relative tolerance is 100.
I tried lower tolerance high absolute tolerance is 1e-10 and the relative tolerance is 1e-5 but the test would fail on TPU.

I am not sure this solution is good though. What is recommended to handle the lower numerical precision on TPU?

[mattjj]

We can just skip the test on the TPU; look for the jtu.skip_on_devices helper.

[mattjj]

I'll merge, then mark this test as skipped on TPU.

[mattjj]

I confirmed internal tests pass after merging this (and after cc53aa9).