Is it possible to save intermediate quantities on disk in a gradient calculation?

[fishjojo]

Dear developers,

In the following example, a and b are kept in memory during the entire calculation. However, they are only needed at the very end of the backward pass. A memory efficient implementation would be to save them on disk after c = a * b, and to read them back from the disk whenever needed in the backward pass. Is it possible to realize this with jax? Or is it possible to define an array class similar to an H5 dataset while it is still traceable by jax? Thank you in advance.

def foo(a, b):
    c = a * b
    return fun(c)

g = jax.grad(foo, 0) (a, b)