Document Haiku version of JAX transforms (hk.jit, ...)

[google-deepmind]

Document Haiku version of JAX transforms (hk.jit, ...)

[sjmielke]

As it's not documented I'm not sure if this is a bug or (not) working as intended:

jnp.exp(0)  # 1.0
jax.jit(jnp.exp)(0)  # 1.0
hk.jit(jnp.exp)(0)  # IndexError: deque index out of range

Lmk if I should delete this comment and report it instead :)

[ibab]

@sjmielke: Thanks for finding that! It fails because the hk.* transforms assume that they run inside of hk.transform right now, but there's no reason why they shouldn't work outside. I have a fix for this in #17.

[trevorcai]

@sjmielke We anticipate that the situations in which you'd want to use hk.jit or hk.grad are limited! They exist as power user workarounds for a particular use cases. I regard them as much more alpha than the rest of Haiku.

Haiku presents a function, hk.transform, which converts impure, object-oriented code with magic functions like hk.get_parameter into JAX transform friendly pure functions.
However, sometimes we want to use JAX transformations inside of a monolithic chunk of Haiku code.
hk.{grad,jit,remat} are exposed for these cases; in all other situations, prefer the JAX equivalent.

For hk.jit:

• We recommend against jitting functions inside of hk.transform that create or use parameters.
  • If you can jax.jit the entire hk.transform(my_fn), do that!
  • If you can't, prefer to extract pure functions representing the expensive portions of your computation and JIT those.
  • In the rare case in which neither of these are possible (e.g. data-dependent control flow that is hard to express with JAX XLA control flow tools, and is hard to break down from a code perspective), then we provide hk.jit as a workaround.

For hk.grad:

• If your model involves taking derivatives inside of your neural network, use hk.grad.

TODO: Add documentation for this stuff. Contributions welcome!

[trevorcai]

We've made a start here with these two:
https://dm-haiku.readthedocs.io/en/latest/transforms.html
https://dm-haiku.readthedocs.io/en/latest/api.html#jax-transforms

Closing for now.