Fixing module imports

graphcore-research · Nov 8, 2023 · 3f9d5ae · 3f9d5ae
1 parent 404a2f6
commit 3f9d5ae
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/jax_scaled_arithmetics/lax/scaled_ops.py b/jax_scaled_arithmetics/lax/scaled_ops.py
@@ -1,5 +1,6 @@
-from core import ScaledArray
-import core
+from jax_scaled_arithmetics import core
+from jax_scaled_arithmetics.core import ScaledArray
+
 from jax import lax
 from functools import partial
 

diff --git a/tests/interpreters/test_interpreter.py b/tests/interpreters/test_interpreter.py
@@ -5,7 +5,7 @@
 import jax.numpy as jnp
 
 from jax_scaled_arithmetics.core import ScaledArray, autoscale
-
+import jax_scaled_arithmetics.lax
 
 class AutoScaleInterpreterTests(chex.TestCase):
  def test__identity(self):