From 6d8f130be7d5e46d31463aa129952104fbcf49dd Mon Sep 17 00:00:00 2001 From: Wassim KABALAN Date: Tue, 9 Jul 2024 02:34:40 +0200 Subject: [PATCH] Import at the start to register everything --- jaxpm/__init__.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/jaxpm/__init__.py b/jaxpm/__init__.py index e69de29..2598fa4 100644 --- a/jaxpm/__init__.py +++ b/jaxpm/__init__.py @@ -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']