Allow using the hk.* version of JAX transforms outside of hk.transform

[ibab]

Right now the hk version of the JAX transforms fail when used outside of hk.transform with IndexError: deque index out of range.
This PR changes the transforms to fall back to the Jax versions when not used with Haiku models.

Alternatively we would throw a clean error saying that users should use jax.* instead.

[trevorcai]

I would prefer if we threw an error saying to use the JAX versions of these transforms instead.

[ibab]

Done.