new_user55

You can write the cuda kernels and directly differentiate them using Enzyme. Last time I checked KOKKOS was still WIP: https://github.com/EnzymeAD/Enzyme/issues?q=kokkos

Really nice work. I love Enzyme, and used it in my project about differentiable atomic descriptors. Idea was that I can quickly gobble up existing C++ and fortran codes alike for atomic descriptors and create a encompassing library what differentiate against hyper-params as well! But at time Enzyme was very early ~0.0.50 version or so. In our observations also Enzyme was fast enough that performance wise it matched the analytical gradients (when embedded inside entire pipeline)  .

I will guess for linux. Most out of the box linux laptops I saw were intel based. I guess open source support of intel is best among others in the industry. Even in my current thinkpad first thing I did was to replace its wifi module from realtek to intel (realtek was always hanging/dropping connection etc).

Its really cool! If you don't mind me asking, does it support variable size inputs? I am bit confused about JAX in that regards. I am trying for long to run JAX stablehlo models in C++ for inference. However dynamic shapes were still an issue. If I understand correctly, it recompiles the kernels for all different shapes at runtime, so if inputs vary too much in shape, it will spend considerable time recompiling kerenels. In C++ inference it becomes impossible .However I could be wrong (I did not fully understand the issue, the developer of gopjrt tried to explain it to me!). Do you have any thoughts on this?

e.g.:

https://github.com/openxla/xla/issues/33092 https://github.com/openxla/xla/issues/35556

Explanation from the gopjrt dev:

https://github.com/gomlx/gopjrt/issues/59

We wanted to use ONNX runtime for a "model driver" for MD simulations, where any ML model can be used for molecular dynamics simulations. Problem was it was way too immature. Like ceiling function will only work with single precision in ONNX. But the biggest issue was that we could not take derivatives in ONNX runtime, so any complicated model that uses derivatives inside was a nogo, is that limitation still exist? Do you know if it can take derivatives in training mode now?

Eventually we went with pytorch only support for the time being, with still exploring OpenXLA in place of ONNX, as a universal adapter: https://github.com/ipcamit/colabfit-model-driver

Tinygrad is python only right? Can it provide gradients during C++ runtime as well? ONNX runtime have multiple language backends for inference.

There is [torch-md](https://github.com/torchmd/torchmd)