Neil Girdhar

Thanks for taking a look! > I prefer not to add pass any additional arguments to nvim by default, and let the user decide if they would like to add...

@mikesmithgh Thanks very much for looking into this, and for providing a perfect fix! Yes, you're exactly right. I'm using auto-session. It seems like you've reproduced my exact problem, and...

@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...

@mtthss Would you mind taking a look at this?

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

> which arguments are causing errors to you? All of the ones I changed. I maintain a [shim library](https://github.com/NeilGirdhar/tjax/tree/main/tjax/_src/gradient) so that I can use optax with dynamic, inspectable parameters. What...

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

> or even better, by marking optimisers as pytrees and splitting 'dynamic' data such as learning rate and such from static data... Exactly! This is the approach I took in...

Just some thoughts: > JAX function transformations on classes can be a source of confusion for users. Functions are generally easier to reason about under JAX transformations than class methods....

Would it be possible to make this work with the Array API? ```python xp = get_namespace(x) return xp.logaddexp(0, x, **kwargs) ```