SYNTOLOGY HomeExplorerAtlasCodeMethodologyAboutDevelopersFeedPricing
Paper · 1912.04232 · 2019

JAX, M.D.: A Framework for Differentiable Physics

arXiv · PDF · Open in the Atlas

Code that ran

We lifted 5 functions out of this paper's own repositories and ran 5 of them in a sandbox. "Ran" means the function executed on a synthesized input and returned a value. It is not a reproduction of the paper's results.

RepositoryRoleRan
google/jax-md canonical 5 of 5
FunctionStatusWhere it lives
canonicalize Ran google/jax-md/jax_md/interpolate.py
code served (permissive licence) · get_code("8daa17a14485429d")
constant Ran google/jax-md/jax_md/interpolate.py
code served (permissive licence) · get_code("2c76bb3ef11ab9af")
dataclass Ran google/jax-md/jax_md/dataclasses.py
code served (permissive licence) · get_code("96b577a29493256c")
estimate_max_neighbors Ran google/jax-md/jax_md/custom_partition.py
code served (permissive licence) · get_code("01b9f68943d5c406")
unpack Ran google/jax-md/jax_md/dataclasses.py
code served (permissive licence) · get_code("a9d0437b7f7e7a53")

Repositories linked to this paper

Some links come from the archived Papers with Code dataset (CC BY-SA 4.0): attribution and licence.

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.

For agents

The same record, over MCP at https://syntology.ai/mcp:

get_harvested_code_for_paper("1912.04232")
get_code_for_paper("1912.04232")
have("1912.04232")

Connect an agent — have() is free.