ΦFlow: Differentiable Simulations for PyTorch, TensorFlow and Jax

Differentiable processes have proven an invaluable tool for machine learning (ML) in scientific and engineering settings, but most ML libraries are not primarily designed for such applications. We present Φ Flow , a Python toolkit that seamlessly integrates with PyTorch, Tensor-Flow, Jax and NumPy, simplifying the process of writing differentiable simulation code at every step. Φ Flow provides many essential features that go beyond the capabilities of the base libraries, such as differential operators, boundary conditions, the ability to write dimensionalityagnostic code, floating-point precision management, fully differentiable preconditioned (sparse) linear solves, automatic matrix generation via function tracing, integration of SciPy optimizers, simulation vectorization, and visualization tools. At the same time, Φ Flow inherits all important traits of the base ML libraries, such as GPU / TPU support, just-in-time compilation, and automatic differentiation. Put together, these features drastically simplify scientific code like PDE or ODE solvers on grids or unstructured meshes, and Φ Flow even includes out-of-the-box support for fluid simulations. Φ Flow has been used in various publications and as a ground-truth solver in multiple scientific data sets. It is available at https: //github.com/tum-pbs/PhiFlow .