Custom PyTree Backend -> Equinox Backend #110
Replies: 4 comments 1 reply
-
Hello @jejjohnson ! Very sorry for missing this, but I'm taking some time away from academia so I was not very active on the repo. I believe that integrating equinox (and diffrax!) is a great idea, do you mind if we have a more detailed conversation when I'm back on this? In general, however, this project is LGPL exactly because I'm more than happy for people to contribute to it or make their own versions / forks. |
Beta Was this translation helpful? Give feedback.
-
Also, regarding the possibility of transforming classes to flax: you can always transform a domain = Domain(...)
def func(x: jnp.ndarray):
u = FourierSeries(params = x, domain=domain)
b = some_jaxdf_operator(u)
return b.on_grid This should work with every For example, in this class we are using |
Beta Was this translation helpful? Give feedback.
-
Thank you for your above comments! I would love to link up sometime in June if you have some time. I've been working silently on trying to improve my understanding of PDEs and how one can use Jax. I did a pure functional approach with equinox and explicit finite difference schemes using finitediffX (and kernex) and timesteppers using diffrax, e.g. see my attempt at the 12 steps to Navier-Stokes. I have to say the combination of libraries made working with PDEs a breeze. My ultimate goal is to work with simple ocean-like PDEs to work with some parameterisations but I think everything I've learned along the way could be very useful for anyone working with PDEs. I would like to advertise this a bit to my numerical PDE peers so that they can see how we can jax-ify all the things. Of course the next step is the jaxdf implicit/automatic way of doing the discretisation. So, again, I would love to link up with you and perhaps discuss ways of doing this. I think the way jaxdf does it is the most intuitive I've seen and it also mirrors a recent example in the diffrax framework for the heat equation. |
Beta Was this translation helpful? Give feedback.
-
Hi! I'm finally back full time on this, sorry for the slow response time so far. Shall we have a chat about this in more details? I made a discord server for this project (you can see it in the README) or I can send you an email and organize a call. |
Beta Was this translation helpful? Give feedback.
-
Hi,
I really like this library a lot! It is the cleanest API for PDEs that I have seen that's intuitive and fairly extendible! I am working with PDEs for oceanography applications like inverse problems and surrogates and I would like to use this API to make some PDEs for Quasi-Geostrophic equations and Shallow Water equations.
I was walking through the library bit by bit to understand how you all built this to see if it would work for my use-case. There are 2 things I wanted to ask if you had an example:
diffrax
. It wasn't updated in the docs but I saw the latest release and code-base. Do you have an example of how to use this with diffrax? I think training the network parameters should be straight forward. However, I'm not sure how it works with mixing and matching different definitions ofpytrees
for the pytrees defined inequinox
and the pytrees define in this library. (But perhaps I am wrong about this).equinox
orflax
. If you have some examples already, that would be great to see! (Again, perhaps it's easier than I think it is...)Also in relation to those two questions, the crux thing I wanted to ask was if you had any intention of converting/upgrading the current framework to use an already implemented library for
pytrees
, e.g.equinox
? For me it seems like a good choice especially to couple this with the already "complete" ODE librarydiffrax
. In fact, a new example was just posted for how one could solve PDEs usingdiffrax
. But I wanted to know the motivations and future directions for the package. I suppose there are 3 potential options below:Integrate
equinox
into thisjaxdf
library as the back-end to make everything compatible with thediffrax
library. I would definitely love to help with this. The downside of course is that now it is no longer agnostic and it might be difficult to mix and match libraries.with your permission of course, I could also do a small rewrite of the library using the
equinox
backend for the pytrees. I could rename the libraryPDiffrax
oreqx-df
. I personally wouldn't mind doing it because it is good for my use case and it would be an excellent exercise for me to become more familiar with JAX.If there is a clever way to "transform" the pytree classes here to other frameworks, e.g.
equinox
,flax
, then that would probably be the best case! (perhaps you already have examples of this?)Beta Was this translation helpful? Give feedback.
All reactions