DeprecationWarning transform_with_state because of jax.xla vs jax.interpreters.xla

[joeryjoery]

Hi, when calling hk.transform_with_state internally there is an access to jax.xla, this is marked as deprecated in favor of jax.interpreters.xla. The problem is in checking if the provided function f to the haiku function is not jax transformed.

I.e., the misdoer is: check_not_jax_transformed (at least, that's how far I've looked into the code; there may exist more references).

Could this be updated? That will help silencing my test output :).