fixed bug in lazy initialization

Author: Alexander

mpx/_dtypes.py

@@ -3,13 +3,11 @@
 import sys
 import types
 
+HALF_PRECISION_DATATYPE = jnp.float16  # Default half precision datatype
+
 # We do to avoid that jax is directly called when importing this module.
-# This is to ensure that mpx works with distributed training.
-class _MaxConstantsLazyInit(types.ModuleType):
-    @property
-    def HALF_PRECISION_DATATYPE(self):
-        return jnp.float16
-    
+# This is to ensure that the constants are lazily initialized.
+class _MaxConstantsLazyInit(types.ModuleType):    
     @property
     def FLOAT16_MAX(self):
         return jnp.ones([], dtype=jnp.float32) * (2 - 2**(-10)) * 2**15

tests/test_dtypes.py

@@ -3,6 +3,7 @@
 from mpx import half_precision_datatype, set_half_precision_datatype
 
 class TestDtypes(unittest.TestCase):
+
     def test_default_half_precision(self):
         """Test that the default half precision datatype is float16"""
         self.assertEqual(half_precision_datatype(), jnp.float16)