feat: import lax argument classes (#99)

Signed-off-by: nstarman <nstarman@users.noreply.github.com>
GalacticDynamics · Nov 23, 2024 · 15bda5d · 15bda5d
1 parent 8ec0ff0
commit 15bda5d
Showing 1 changed file with 30 additions and 0 deletions.
diff --git a/src/quaxed/lax/__init__.py b/src/quaxed/lax/__init__.py
@@ -172,6 +172,19 @@
     "with_sharding_constraint",
     # ----- Linear Algebra Operators -----
     "linalg",
+    # ----- Argument classes -----
+    "ConvDimensionNumbers",
+    "ConvGeneralDilatedDimensionNumbers",
+    "DotAlgorithm",
+    "DotAlgorithmPreset",
+    "FftType",
+    "GatherDimensionNumbers",
+    "GatherScatterMode",
+    "Precision",
+    "PrecisionLike",
+    "RandomAlgorithm",
+    "RoundingMethod",
+    "ScatterDimensionNumbers",
 ]
 
 
@@ -184,6 +197,23 @@
 
 from . import linalg
 
+# Explicit imports that don't need to be quaxified
+# isort: split
+from jax.lax import (
+    ConvDimensionNumbers,
+    ConvGeneralDilatedDimensionNumbers,
+    DotAlgorithm,
+    DotAlgorithmPreset,
+    FftType,
+    GatherDimensionNumbers,
+    GatherScatterMode,
+    Precision,
+    PrecisionLike,
+    RandomAlgorithm,
+    RoundingMethod,
+    ScatterDimensionNumbers,
+)
+
 
 def __dir__() -> list[str]:
     """List the module contents."""