Repository: google-deepmind/optax

Issue:
`_bias_correction` forces recompilation of `integer_pow`
This might be related to #197, but I don't have enough knowledge to confirm it.

In short, while using `optax` with `flax.linen`, `optimizer.update` forces jit-recompilation of the `integer_pow` jax primitive.
I get these results on a GPU backend.

```python
from functools import partial
import jax
import jax.numpy as jnp
import optax
from flax import linen as nn
from jax.config import config
import os


config.update("jax_log_compiles", 1)
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"


def biascorrection_recompilation_test():
    class M(nn.Module):
        @nn.compact
        def __call__(self, x):
            x = nn.Dense(1)(x)
            return x

    m = M()
    o = optax.adam(0.025, b1=int(1))

    input_array = jnp.ones((8,), dtype=float)
    params = m.init(jax.random.PRNGKey(0), input_array)
    opt_state = o.init(params)

    m_forward = jax.vmap(m.apply, in_axes=(None, 0))

    @jax.jit
    @partial(jax.value_and_grad, has_aux=True)
    def loss_fn(params, x, y):
        params_m = params
        y_m = m_forward(params_m, x)
        loss = jnp.mean(y_m - y ** 2)
        return loss, ()



    for i in range(1000):
        print(i)
        inputs = jax.random.normal(jax.random.PRNGKey(i), (32, 8))
        outputs = jax.random.normal(jax.random.PRNGKey(-i), (32, 8))
        (loss, aux), grads = loss_fn(params, inputs, outputs)
        updates, opt_state = o.update(grads, opt_state)
        params = optax.apply_updates(params, updates)


if __name__ == "__main__":
    biascorrection_recompilation_test()

```
At each `i`:
```
WARNING:absl:Finished tracing + transforming prim_fun for jit in 0.00016188621520996094 sec
WARNING:absl:Compiling prim_fun (140402616089792 for args (ShapedArray(float32[], weak_type=True),).
WARNING:absl:Finished XLA compilation of integer_pow in 0.02527928352355957 sec
```

Where `prim_fun` refers the `integer_pow` op at this line:
https://github.com/deepmind/optax/blob/b4aa6657bbf79985279dea76eaf6d53b25d7e8d9/optax/_src/transform.py#L105

#329 does remove the issue by replacing the call to`integer_pow` with `pow`, but I am not sure this is the direction you would like to take.


Discussion:
Hello @epignatelli,

I am not totally sure why the recompilation is happening on every step (for integer pow) - although, this should not happen if you jit the the full update function, and apply that jitted function in every loop iteration.

e.g.

```python
def update_fn(params, inputs, outputs)

    @partial(jax.value_and_grad, has_aux=True)
    def loss_fn(params, x, y):
        params_m = params
        y_m = m_forward(params_m, x)
        loss = jnp.mean(y_m - y ** 2)
        return loss, ()

    (loss, aux), grads = loss_fn(params, inputs, outputs)
    updates, opt_state = o.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    return params


update_fn = jax.jit(update_fn)

for i in range(1000):
    print(i)
    inputs = jax.random.normal(jax.random.PRNGKey(i), (32, 8))
    outputs = jax.random.normal(jax.random.PRNGKey(-i), (32, 8))
    params = update_fn(params, inputs, outputs)
```



@hbq1 suggests that the recompilation for interger_pow may be happening because the exponent is a static argument for integer pow. This would be because more efficient algorithms exist when the exponent is a known integer.

Does the above help?

It would be interesting to see what happens when jitting the whole function as above - maybe a different algorithm would be selected. :) 
Indeed, [the docs](https://github.com/google/jax/blob/main/jax/_src/lax/lax.py#L362) state that the exponent in `integer_pow` is a _fixed_ integer.
(I just checked, and this is not related to #197 as much as I wish it were.)

Resolving patch:
diff --git a/optax/_src/transform.py b/optax/_src/transform.py
--- a/optax/_src/transform.py
+++ b/optax/_src/transform.py
@@ -109,9 +109,8 @@ def orderth_norm(g):
 
 def bias_correction(moment, decay, count):
   """Perform bias correction. This becomes a no-op as count goes to infinity."""
-  bias_correction_ = 1 - decay**count
   return jax.tree_util.tree_map(
-      lambda t: t / bias_correction_.astype(t.dtype), moment)
+      lambda t: t / (1 - decay ** count.astype(t.dtype)), moment)
 
 
 def _reject_complex(params):


Test patch: