guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-04-24 19:50:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Import at the start to register everything

Browse source
This commit is contained in:
Wassim KABALAN 2024-07-09 02:34:40 +02:00
parent e506dc2a7e
commit 6d8f130be7
1 changed files with 8 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
jaxpm/__init__.py
View file

@ -0,0 +1,8 @@
# Execute the register_operator functions
import jaxpm._src.base_ops
import jaxpm._src.painting_ops
import jaxpm.ops
import jaxpm.painting
from jaxpm._src.spmd_config import SPMDConfig
__all__ = ['SPMDConfig']