JAX-Fluids
简介
JAX-Fluids 是一个基于 Google JAX 框架的全自动可微 CFD（计算流体力学）求解器，专注于端到端可微的流体仿真。相比传统 CFD 求解器（如 OpenFOAM），它允许对仿真参数（几何、边界条件、物性）做梯度反传，因此天然契合物理信息机器学习（Physics-Informed ML）、优化设计与神经网络耦合仿真。
核心特性
- 全自动可微：对所有仿真参数可求导，无需手动伴随方程
- 多GPU支持：JAX 的 SPMD 抽象天然支持多 GPU/TPU 并行，数据并行零改动
- 高阶精度：WENO-JS 格式，支持 3rd~5th 阶精度
- 多相流：支持两相流（VOF 方法）
- 与神经网络耦合：可作为 Neural ODE / PINN 的物理前向求解器
适用场景
- 形状优化（机翼、管道、钝体）：梯度直接反传至设计参数
- 物理信息神经网络（PINN / PIAN）前向求解器
- 数据驱动湍流建模（closure term 学习）
- 端到端物理仿真-控制联合学习
使用注意
- License 标注 MIT，但原始论文 CPC 部分标注 GPLv3，建议仅作算法参考不直接嵌入闭源商业产品
- 适合学术研究；工业级验证建议与传统求解器（OpenFOAM/ANSYS）交叉验证
- 国内访问 GitHub 需代理，可配合 Gitee 镜像使用
标签: JAX可微CFD可微物理GPU端到端
相关条目
- JAX-CFD · Google Research 基于 JAX 的可微流体求解器，支持端到端梯度反传，用于把神经网络嵌入 CFD 时间推进、做可微物理与混合建模研究。
- 使用 Agentic AI 将 CUDA Tile 操作从 Python 翻译到 Rust · cuTile Rust (cutile-rs) 是用于在 Rust 中安全编写 GPU 内核的瓦片化系统，扩展 Rust 所有权模型实现内存安全。
- NVIDIA Transformer Engine 在 JAX 中加速无Drop MoE训练 · NVIDIA Transformer Engine 新增JAX支持的MoE优化内核，助力大规模AI模型高效稀疏训练。
- DeepXDE v1.15.0 · DeepXDE v1.15.0 发布，新增 L-LAAF 激活函数与 PyTorch 后端 MPS GPU 支持，提升微分方程深度学习求解效率。
- AIPerf：大规模LLM推理性能基准测试工具 · AIPerf 是 NVIDIA 推出的 LLM 推理基准测试工具，用于在部署场景下评估大规模语言模型的推理速度和性能表现。