Allow gradient transform parameters to be dynamic

[NeilGirdhar]

No description provided.

[NeilGirdhar]

@hawkinsp Pinging you since you recently repaired some type annotation errors. The optimizer classes accepting only float breaks type annotations for the Tjax shim classes (https://github.com/NeilGirdhar/tjax/blob/main/tjax/_src/gradient/transforms.py). Tjax provides a parallel set of optimizers, identical in functionality, except they support dynamic optimizer parameters. They do this by storing dynamic fields in a dataclass rather than closing over parameters.

However, the optimizer functionality is delegated to Optax, which means calling Optax update methods with Jax arrays. Is there any reason Optax methods can't accept such arrays? Would it be possible to widen these parameter types to jax.Array | float?

[NeilGirdhar]

@mtthss Would you mind taking a look at this?

[mtthss]

Hello. I was on paternity leave for most of the past year. Are you still having this issue? Happy to look into it if that's the case

[NeilGirdhar]

@mtthss Hello, yes I'm still getting the type errors. (Congrats on becoming a father!)

[mtthss]

which arguments are causing errors to you?

[NeilGirdhar]

which arguments are causing errors to you?

All of the ones I changed. I maintain a shim library so that I can use optax with dynamic, inspectable parameters. What I ended up doing for the time being is to mark every use of optax with pyright: ignore.

Thanks for taking a look at this.

[NeilGirdhar]

(Of course, my dream would be that you adopt the dynamic design so that I don't have to maintain my shim library 😄.)

[NeilGirdhar]

@mtthss Do you mind taking a look at this?

[rdyro]

I'm a new maintainer. At some point, optax moved to https://optax.readthedocs.io/en/latest/api/optimizer_schedules.html#optax.schedules.inject_hyperparams to make parameters available in the state.

I think the point of these type annotations is to explicitly indicate that (like you're pointing out) the parameters are closed over and not really modifiable without inject_hyperparams.

[NeilGirdhar]

https://optax.readthedocs.io/en/latest/api/optimizer_schedules.html#optax.schedules.inject_hyperparams to make parameters available in the state.

That's a separate issue from this PR. (I believe, in fact, it's the first issue #1.)

I think the point of these type annotations is to explicitly indicate that (like you're pointing out) the parameters are closed over and not really modifiable without inject_hyperparams.

Even if the parameters are closed over, that doesn't make them floats. A closed-over Jax array is still a Jax array. The annotations are simply wrong. It's just that most users probably don't hit this because most users are probably just using fixed float values.

[rdyro]

Let me talk to @vroulet

[vroulet]

That's a completely fair point. We'll make an attempt to change the types (this will affect all internal code so many bugs possible) but a priori it should not pose problems. We'll keep you updated Neil. (It'll be next month a priori)