dionhaefner

We thought it would be cool to be able to run GPU workloads across process boundaries from JAX and torch. Train your torch networks through Julia, C++ solvers with adjoints, or any arcane autodiff potion you choose to concoct – now certified fresh without host round trips. Bon apetit!

No idea lol. I assume it’s possible since both Enzyme and GPU programming are pervasive in Julia. Let us know if you end up trying.

Happy to take the blame for the lack of persuasian and engagement then :) Thanks for the feedback, although I’d like to believe there’s a way to have that cake and eat it too.

Indeed, just like I let my compiler write (at least) chunks of my AD logic. Not great when the tool becomes a leaky abstraction, but overall net positive don't you think?

Author here — I work on Tesseract at Pasteur Labs, and I wrote this up because the "what if this was possible" was bugging me for way too long :)

I was surprised by how well this worked, the LFortran + Enzyme stack seems to be a very clean way to get gradients through Fortran code via LLVM IR transformations. Pretty cool to see a 220-line Fortran heat solver turn into ~6,900-line reverse pass automatically if I dare say so.

Would be awesome to see this applied to a real scientific codebase, and I hope that the demo is enough to convince people that it’s worth trying.

This is a case study in making gradient-based optimization (a la gradient descent) work with tools that weren’t designed for it. The goal is optimizing rocket grid fin stiffness (8 bars, 16 angular parameters), with a pipeline that includes Ansys SpaceClaim (CAD), PyMAPDL (FEM solver), and JAX. Each step needs derivatives, which means stitching together analytical adjoints from Ansys, finite differences for mesh operations, and JAX autodiff for everything else. What makes this practical is that each component runs in its own isolated environment (Tesseract), so the Ansys tools can live on Windows with their license requirements while the optimization logic runs on Linux while still being composable for end-to-end autodiff.

The payoff is watching emergent behavior: starting from random bar positions, the optimizer discovers that orthogonal grid patterns work well, diagonal lateral bars create efficient load paths, and clustering material near the attachment points maximizes stiffness. Final result is 75% stiffer than random and 24% better than a regular grid, even after adding back symmetry constraints for manufacturing.

Climate models are vastly complex, and you need to bring together many experts from many disciplines to write and maintain one, and analyze the output. This seems to lead to the simplest methods coming out on top. Perhaps it could be solved with better abstractions (a lot of very smart people are trying).

I agree with your point regarding non-composability and ”algorithm lock in” (which may or may not be solvable woth better abstractions), but explicit time stepping schemes are still the main workhorse of global ocean modelling, so I’m not sure whether ”silly” is the right label here.

The thing with reduced precision is that things may look fine at first, but then you eventually notice unphysical features in your solution (like additional wave modes after very long simulation times, or energy conservation issues). So we really don't know as a community yet how far we can venture from float64, but it looks like float32 may be viable.

Veros works OK on TPUs (about the same speed as a high-end GPU), but since you can't buy TPUs that's an immediate no for most academic users of climate models. Renting hardware doesn't really make sense when you keep it busy for months at a time and the HPC infrastructure is already in place.

True, but unfortunately Pytorch is not quite there yet when it comes to more complex benchmarks:

https://github.com/dionhaefner/pyhpc-benchmarks#example-resu...

JAX really is the only library that comes close to low-level code on CPU, almost always (that I've tried).

I was mostly referring to the millions (billions?) of dollars getting poured into Python library development by tech companies. With the effect that Python stays relevant and has a thriving library ecosystem. Maybe I'm wrong and Julia is just that good that it doesn't matter - I guess time will tell.

A bit late to the party, but here are some reasons:

- When we started Veros (~4 years ago) Julia was very new on our radar and we didn't know whether it would stick. And to be frank, I'm still not convinced whether it will stick. Yes it seems like a fantastic language, but we all know how long it took Python to gain traction.

- Climate scientists and students already do their post-processing in Python. Having the whole stack in the same language makes things a lot easier for domain experts whose first priority is physics, not coding.

- Python skills translate better to other jobs, which I think is important for young academics.

- The Python library ecosystem is so good. Need to use PETSc? `import petsc4py`. Simplify postprocessing? Export your model state as `xarray` dataset. Julia is great for bleeding edge autodiff through everything stuff, but the bread and butter libraries are just so polished and battle tested in Python.

- I don't know Julia :)

Yes that would probably work! GAMs are a bit of a black box to me so I find it hard to reason about them. I think my method is a bit simpler and it’s more straightforward to get uncertainties, but certainly less elegant than a smooth solution.

I basically determine p via Bayesian inference within every bin (via a conjugate beta prior for p which gives a beta posterior). If that’s not Bayesian then I don’t know what is :)

Yes the pruning can be done with a frequentist method too. Yes you can come up with smarter / more statistically sound ways to construct these binnings. Do they work on >1e9 data points?

Yes, it’s possible. There are some proofs of concept for that IIRC. If you prefer to write everything in Javascript you can do that and do all heavy lifting in the client (plus some extra work to provide indexing and metadata like Terracotta).

Only if the client has access to the full bit depth. The rendered PNGs are uint8, but the underlying raw data is often stored in 16 or even 32 bit precision.

I’m not sure which standards you are referring to (I guess that means no). If you have a link or open an issue we could scope this out.

Rasterio wheels ship with GDAL binaries. No C compiler needed.

That is true to some degree, but with Terracotta we allow the frontend to do manipulations on the fly. Like changing contrast, applying colormaps, or band math. So the number of possible tiles is infinite.

If you don’t need that, you can just pre-generate all PNGs - then you don’t even need a server anymore, just a hash function to identify the right file and a big disk.

I think the point was that Terracotta might be better off using GDAL directly instead of via rasterio.

Maybe, but I think mapbox is doing a great job with rasterio, and so far we didn’t run into any unsolvable problems (although we got close sometimes). I think for the ease of installation alone rasterio is a blessing.

Author here. Not really, since we rely on cloud-optimized GeoTiffs for performance, in particular their overviews. This makes it so we can serve tiles from huge datasets (10s of GB) in less than 100ms. I don’t think that would be possible with NetCDF.

Yes. Fortran is used extensively in climate modelling (even newer models).