{"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-md-end-to-end-differentiable-hardware-1","title":"JAX, M.D.: A Framework for Differentiable Physics","arxiv_id":"1912.04232","date":"2019-12-09","proceeding":null,"authors":["Samuel S. Schoenholz","Ekin D. Cubuk"],"abstract":"We introduce JAX MD, a software package for performing differentiable physics simulations with a focus on molecular dynamics. JAX MD includes a number of physics simulation environments, as well as interaction potentials and neural networks that can be integrated into these environments without writing any additional code. Since the simulations themselves are differentiable functions, entire trajectories can be differentiated to perform meta-optimization. These features are built on primitive operations, such as spatial partitioning, that allow simulations to scale to hundreds-of-thousands of particles on a single GPU. These primitives are flexible enough that they can be used to scale up workloads outside of molecular dynamics. We present several examples that highlight the features of JAX MD including: integration of graph neural networks into traditional simulations, meta-optimization through minimization of particle packings, and a multi-agent flocking simulation. JAX MD is available at www.github.com/google/jax-md.","url_abs":"https://arxiv.org/abs/1912.04232v2","url_pdf":"https://arxiv.org/pdf/1912.04232v2.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-md-end-to-end-differentiable-hardware-1","repo_url":"https://github.com/google/jax-md","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"Apache-2.0"}}],"tasks":[{"task_slug":"drug-discovery","task_name":"Drug Discovery"},{"task_slug":null,"task_name":"GPU"}],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"syntology_url":"https://syntology.ai/paper/1912.04232","atlas_url":"https://app.syntology.ai/?focus=1912.04232","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"1912.04232"}},"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/google/jax-md","reach":{"status":"ok","spdx":"Apache-2.0"}}],"summary":{"unverified":5},"by_repo_kind":{"official":{"samples":5,"ran":0,"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":0,"samples":[{"code_sha256_prefix":"8daa17a14485429d","entry":"canonicalize","repo":"google/jax-md","repo_kind":"official","path":"jax_md/interpolate.py","file_url":"https://github.com/google/jax-md/blob/HEAD/jax_md/interpolate.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"8daa17a14485429d"}},{"code_sha256_prefix":"2c76bb3ef11ab9af","entry":"constant","repo":"google/jax-md","repo_kind":"official","path":"jax_md/interpolate.py","file_url":"https://github.com/google/jax-md/blob/HEAD/jax_md/interpolate.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"2c76bb3ef11ab9af"}},{"code_sha256_prefix":"96b577a29493256c","entry":"dataclass","repo":"google/jax-md","repo_kind":"official","path":"jax_md/dataclasses.py","file_url":"https://github.com/google/jax-md/blob/HEAD/jax_md/dataclasses.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"96b577a29493256c"}},{"code_sha256_prefix":"01b9f68943d5c406","entry":"estimate_max_neighbors","repo":"google/jax-md","repo_kind":"official","path":"jax_md/custom_partition.py","file_url":"https://github.com/google/jax-md/blob/HEAD/jax_md/custom_partition.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"01b9f68943d5c406"}},{"code_sha256_prefix":"a9d0437b7f7e7a53","entry":"unpack","repo":"google/jax-md","repo_kind":"official","path":"jax_md/dataclasses.py","file_url":"https://github.com/google/jax-md/blob/HEAD/jax_md/dataclasses.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"a9d0437b7f7e7a53"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}