{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/jax-fluids-2-0-towards-hpc-for-differentiable","title":"JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows","arxiv_id":"2402.05193","date":"2024-02-07","proceeding":null,"authors":["Deniz A. Bezgin","Aaron B. Buhendwa","Nikolaus A. Adams"],"abstract":"In our effort to facilitate machine learning-assisted computational fluid dynamics (CFD), we introduce the second iteration of JAX-Fluids. JAX-Fluids is a Python-based fully-differentiable CFD solver designed for compressible single- and two-phase flows. In this work, the first version is extended to incorporate high-performance computing (HPC) capabilities. We introduce a parallelization strategy utilizing JAX primitive operations that scales efficiently on GPU (up to 512 NVIDIA A100 graphics cards) and TPU (up to 1024 TPU v3 cores) HPC systems. We further demonstrate the stable parallel computation of automatic differentiation gradients across extended integration trajectories. The new code version offers enhanced two-phase flow modeling capabilities. In particular, a five-equation diffuse-interface model is incorporated which complements the level-set sharp-interface model. Additional algorithmic improvements include positivity-preserving limiters for increased robustness, support for stretched Cartesian meshes, refactored I/O handling, comprehensive post-processing routines, and an updated list of state-of-the-art high-order numerical discretization schemes. We verify newly added numerical models by showcasing simulation results for single- and two-phase flows, including turbulent boundary layer and channel flows, air-helium shock bubble interactions, and air-water shock drop interactions.","url_abs":"https://arxiv.org/abs/2402.05193v1","url_pdf":"https://arxiv.org/pdf/2402.05193v1.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"jax-fluids-2-0-towards-hpc-for-differentiable","repo_url":"https://github.com/tumaer/jaxfluids","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"jax","reach":null},{"paper_slug":"jax-fluids-2-0-towards-hpc-for-differentiable","repo_url":"https://github.com/WangWen-kang/JAX-BC","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"jax","reach":null}],"tasks":[{"task_slug":null,"task_name":"GPU"}],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"syntology_url":"https://syntology.ai/paper/2402.05193","atlas_url":"https://app.syntology.ai/?focus=2402.05193","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2402.05193"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-25T09:33:49+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/tumaer/jaxfluids","reach":null},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/WangWen-kang/JAX-BC","reach":null}],"summary":{"ran_draft_wrong":2},"by_repo_kind":{"listed":{"samples":1,"ran":1,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":2,"samples":[{"code_sha256_prefix":"2f3c8b0b993d9090","entry":"build_loss_fn","repo":null,"repo_kind":null,"path":null,"file_url":null,"link_basis":"identical_code_first_harvested_elsewhere","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":null,"behaviour_fingerprint":false,"licence":null,"inline_ok":false,"mcp_get_code":{"code_sha256":"2f3c8b0b993d9090"}},{"code_sha256_prefix":"be2b79790f4f8650","entry":"get_face_slices_retrieve","repo":"WangWen-kang/JAX-BC","repo_kind":"listed","path":"jaxfluids/halos/outer/boundary_condition.py","file_url":"https://github.com/WangWen-kang/JAX-BC/blob/HEAD/jaxfluids/halos/outer/boundary_condition.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":"deterministic","behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"be2b79790f4f8650"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}