c0g

Thanks for the answer :-) glad to hear it’s being worked on. Anything early I can play with?

Any update on ETA here? I am trying to use Brax on Metal and it wants the cholesky decomp. @kulinseth

Is jax-metal open source? I can’t find it but would consider contributing.

I can't currently share the model arch (I hate myself for saying that), but it's a fairly basic transformer. To be clear I didn't mean training diverges, rather loading my...