Owen L
Owen L
Definitely should be possible to clean it up a lot, but any preliminary thoughts on this approach would be much appreciated: https://github.com/patrick-kidger/diffrax/pull/402
Possible this is an error in equinox, but I wasn't able to exactly replicate it without diffrax, e.g. works fine ```python import os import multiprocessing os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count={}".format( multiprocessing.cpu_count() )...
Issue that the error code references in stack trace: https://github.com/google/jax/issues/13554
All of the above code works 100% fine in 0.4.26 btw (well the sharding is still slower, but that's a different issue)
Great, closing! 0.4.28 fixed
Before going further (there is a lot I am going to improve/polish) I wanted to check with your thoughts on the general approach of KL being terms and exposing the...
I think your idea makes a lot of sense, and I made a fair amount of progress on the solver wrapper approach.
Ok, I polished things up. I went with a sort of hybrid approach where the users specifies the SDEs as you described, then just wraps a solver and everything works...
Thanks @frankschae , good catch!
How important is Flax to the migration? I've been exploring pgx and it seems really nice and I've started using it, but my go to is equinox (https://github.com/patrick-kidger/equinox, probably the...