文件名称：mpi4jax:JAX阵列的零拷贝MPI通信，用于Python中的涡轮增压HPC应用程序
文件大小：4.56MB
文件格式：ZIP
更新时间：2024-04-25 10:54:05
gpu mpi parallel-computing jit high-performance-computing
mpi4jax mpi4jax支持阵列的零复制，多主机通信，甚至可以通过jitted代码和GPU内存进行通信。 但为什么？ JAX框架，但是其仍然受到限制。 使用mpi4jax ，您可以将基于JAX的模拟扩展到整个CPU和GPU集群（无需离开jax.jit ）。 本着差异化编程的精神， mpi4jax还支持通过一些MPI操作进行差异化。 快速安装 mpi4jax可通过pip和conda ： $ pip install mpi4jax # Pip $ conda install -c conda-forge mpi4jax # conda 我们的文档包括一些更高级的安装示例。 用法示例 from mpi4py import MPI import jax import jax . numpy as jnp import mpi4jax
【文件预览】：
mpi4jax-master
----mpi4jax()
--------_deprecations.py(1KB)
--------_version.py(18KB)
--------__init__.py(797B)
--------_src()
----conf()
--------travis-install-mpi.sh(4KB)
--------ompi_rootenv.patch(1KB)
--------install-cuda-ubuntu.sh(5KB)
----MANIFEST.in(220B)
----README.rst(3KB)
----.gitattributes(47B)
----.flake8(207B)
----.github()
--------workflows()
----versioneer.py(69KB)
----examples()
--------shallow_water.py(16KB)
----setup.cfg(162B)
----LICENSE.md(1KB)
----setup.py(6KB)
----.pre-commit-config.yaml(255B)
----docs()
--------make.bat(795B)
--------shallow-water.rst(4KB)
--------conf.py(2KB)
--------usage.rst(3KB)
--------installation.rst(3KB)
--------shallow-water-source.rst(241B)
--------environment.yml(174B)
--------sharp-bits.rst(4KB)
--------index.rst(235B)
--------Makefile(634B)
--------api.rst(601B)
--------_static()
----tests()
--------test_invalid_jaxlib.py(913B)
--------test_examples.py(527B)
--------test_flush.py(142B)
--------test_deprecations.py(2KB)
--------test_decorators.py(1KB)
--------conftest.py(376B)
--------collective_ops()
--------test_validation.py(2KB)
----.readthedocs.yml(193B)
----.gitignore(2KB)
----pyproject.toml(646B)