JaxPM/jaxpm/__init__.py
2024-07-09 02:34:40 +02:00

8 lines
211 B
Python

# 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']