added unittest

Author: Alexander

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}

mpfj/__init__.py

@@ -1,7 +1,3 @@
-from . import utils
-from . import layers
-from . import optimizers
-
 """
 Mixed Precision for JAX (mpfj)
 
@@ -10,4 +6,4 @@
 
 __version__ = "0.1.0"
 
-from .dtypes import set_half_precision_datatype
+from .dtypes import half_precision_datatype, set_half_precision_datatype

mpfj/cast.py

@@ -34,7 +34,7 @@
 
 from jaxtyping import Array, Float, Int, PyTree, PRNGKeyArray 
 
-from .dtypes import HALF_PRECISION_DATATYPE
+from .dtypes import half_precision_datatype
 
 def cast_tree(tree: PyTree, dtype):
     """
@@ -125,7 +125,7 @@ def cast_to_half_precision(x: PyTree) -> PyTree:
         PyTree: A new PyTree with all elements cast to the half-precision datatype.
     """
     """Cast to half precision (float16/bfloat16, depending on with what we called set_half_precision_datatype)."""
-    return cast_tree(x, HALF_PRECISION_DATATYPE)
+    return cast_tree(x, half_precision_datatype())
 
 
 def force_full_precision(func, return_dtype=jnp.float16):

mpfj/dtypes.py

@@ -9,4 +9,8 @@ def set_half_precision_datatype(datatype):
     Args:
         datatype: The datatype to set as half precision (e.g., jnp.float16).
     """
-    HALF_PRECISION_DATATYPE = datatype
+    global HALF_PRECISION_DATATYPE
+    HALF_PRECISION_DATATYPE = datatype
+
+def half_precision_datatype():
+    return HALF_PRECISION_DATATYPE

mpfj/grad_tools.py

@@ -33,8 +33,8 @@
 
 import optax
 
-import cast as cast
-import loss_scaling as loss_scaling
+import mpfj.cast as cast
+import mpfj.loss_scaling as loss_scaling
 
 from jaxtyping import PyTree, Bool
 
@@ -167,6 +167,8 @@ def optimizer_update(model: PyTree, optimizer: optax.GradientTransformation, opt
     updates, new_optimizer_state = optimizer.update(
         grads, optimizer_state, eqx.filter(model, eqx.is_array)
     )
+    print(updates)
+    print("dddddddd")
     new_model = eqx.apply_updates(model, updates)
 
     # only apply updates to the model and optimizer state if gradients are finite

mpfj/loss_scaling.py

@@ -57,9 +57,9 @@ def all_finite(tree: PyTree) -> Array:
         leaves = map(jnp.isfinite, leaves)
         leaves = map(jnp.all, leaves)
         return jnp.stack(list(leaves)).all()
-    
 
-def scaled(func: callable, scaling: DynamicLossScaling):
+
+def scaled(func: callable, scaling: 'DynamicLossScaling'):
     def wrapper(*_args, **_kwargs):
         value = func(*_args, **_kwargs)
         value = scaling.scale(value)
@@ -93,7 +93,7 @@ def unscale(self, tree):
         inv_loss_scaling = inv_loss_scaling.astype(jnp.float32)   # cast to float32, so the result is float32 (otherwise the whole scaling point would be senseless)
         return jax.tree_util.tree_map(lambda x: x * inv_loss_scaling[0], tree)
 
-    def adjust(self, grads_finite: jnp.ndarray) -> DynamicLossScaling:
+    def adjust(self, grads_finite: jnp.ndarray) -> 'DynamicLossScaling':
         """Returns the next state dependent on whether grads are finite."""
         assert grads_finite.ndim == 0, "Expected boolean scalar"
 
@@ -113,7 +113,7 @@ def adjust(self, grads_finite: jnp.ndarray) -> DynamicLossScaling:
             jnp.maximum(self.min_loss_scaling, self.loss_scaling / self.factor))
 
         # clip to maximum float16 value.
-        loss_scaling = jnp.clip(loss_scaling, a_min=self.min_loss_scaling, a_max=(2 - 2**(-10)) * 2**15)
+        loss_scaling = jnp.clip(loss_scaling, min=self.min_loss_scaling, max=(2 - 2**(-10)) * 2**15)
 
         counter = ((self.counter + 1) % self.period) * grads_finite
 

tests/conftest.py

@@ -0,0 +1,6 @@
+import os
+import sys
+
+# Add the project root directory to the Python path
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, project_root) 

tests/test_cast.py

@@ -0,0 +1,193 @@
+import unittest
+import jax
+import jax.numpy as jnp
+import equinox as eqx
+from jaxtyping import Array, Float, Int, PyTree
+import numpy as np
+
+from mpfj.cast import (
+    cast_tree,
+    cast_to_float32,
+    cast_to_float16,
+    cast_to_bfloat16,
+    cast_to_full_precision,
+    cast_to_half_precision,
+    force_full_precision,
+)
+from mpfj.dtypes import HALF_PRECISION_DATATYPE
+
+
+class EQXModuleBase(eqx.Module):
+    a: Array
+    b: Array
+
+    def __init__(self):
+        self.a = jnp.ones(10, dtype=jnp.float32)
+        self.b = jnp.ones(10, dtype=jnp.float32)
+
+class LeafClass:
+    """If implemented correctly, this class should not be casted"""
+    a: Array
+    b: Array
+
+    def __init__(self):
+        self.a = jnp.ones(10, dtype=jnp.float32)
+        self.b = jnp.ones(10, dtype=jnp.float32)
+
+class EQXModule1(eqx.Module):
+    a: list[EQXModuleBase]
+    b: Array
+    c: LeafClass
+
+    def __init__(self):
+        self.a = [EQXModuleBase() for _ in range(10)]
+        self.b = jnp.ones(10, dtype=jnp.float32)
+        self.c = LeafClass()
+
+
+class TestCastFunctions(unittest.TestCase):
+    def setUp(self):
+        # Create some test data
+        self.array_float32 = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
+        self.array_float16 = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float16)
+        self.array_bfloat16 = jnp.array([1.0, 2.0, 3.0], dtype=jnp.bfloat16)
+        self.nested_dict = {
+            'a': self.array_float32,
+            'b': {'c': self.array_float16, 'd': self.array_bfloat16}
+        }
+        self.mixed_tree = {
+            'array': self.array_float32,
+            'scalar': 42,
+            'nested': {
+                'array': self.array_float16,
+                'none': None
+            }
+        }
+    
+    def test_cast_eqx_module(self):
+        # Create test module
+        module = EQXModule1()
+        
+        # Test casting to float16
+        result = cast_tree(module, jnp.float16)
+        # Check that arrays in nested EQXModuleBase instances are cast
+        for base_module in result.a:
+            self.assertEqual(base_module.a.dtype, jnp.float16)
+            self.assertEqual(base_module.b.dtype, jnp.float16)
+        # Check direct array is cast
+        self.assertEqual(result.b.dtype, jnp.float16)
+        # Check that LeafClass arrays are NOT cast since it's not an eqx.Module
+        self.assertEqual(result.c.a.dtype, jnp.float32)
+        self.assertEqual(result.c.b.dtype, jnp.float32)
+
+        # Test casting to bfloat16 
+        result = cast_tree(module, jnp.bfloat16)
+        # Check nested modules
+        for base_module in result.a:
+            self.assertEqual(base_module.a.dtype, jnp.bfloat16)
+            self.assertEqual(base_module.b.dtype, jnp.bfloat16)
+        self.assertEqual(result.b.dtype, jnp.bfloat16)
+        # LeafClass should remain unchanged
+        self.assertEqual(result.c.a.dtype, jnp.float32)
+        self.assertEqual(result.c.b.dtype, jnp.float32)
+
+        # Test casting back to float32
+        result = cast_tree(module, jnp.float32)
+        for base_module in result.a:
+            self.assertEqual(base_module.a.dtype, jnp.float32)
+            self.assertEqual(base_module.b.dtype, jnp.float32)
+        self.assertEqual(result.b.dtype, jnp.float32)
+        self.assertEqual(result.c.a.dtype, jnp.float32)
+        self.assertEqual(result.c.b.dtype, jnp.float32)
+
+    def test_cast_tree(self):
+        # Test casting to float32
+        result = cast_tree(self.array_float16, jnp.float32)
+        self.assertEqual(result.dtype, jnp.float32)
+        
+        # Test casting nested structure
+        result = cast_tree(self.nested_dict, jnp.float32)
+        self.assertEqual(result['a'].dtype, jnp.float32)
+        self.assertEqual(result['b']['c'].dtype, jnp.float32)
+        self.assertEqual(result['b']['d'].dtype, jnp.float32)
+
+    def test_cast_to_float32(self):
+        result = cast_to_float32(self.array_float16)
+        self.assertEqual(result.dtype, jnp.float32)
+        
+        result = cast_to_float32(self.nested_dict)
+        self.assertEqual(result['a'].dtype, jnp.float32)
+        self.assertEqual(result['b']['c'].dtype, jnp.float32)
+        self.assertEqual(result['b']['d'].dtype, jnp.float32)
+
+    def test_cast_to_float16(self):
+        result = cast_to_float16(self.array_float32)
+        self.assertEqual(result.dtype, jnp.float16)
+        
+        result = cast_to_float16(self.nested_dict)
+        self.assertEqual(result['a'].dtype, jnp.float16)
+        self.assertEqual(result['b']['c'].dtype, jnp.float16)
+        self.assertEqual(result['b']['d'].dtype, jnp.float16)
+
+    def test_cast_to_bfloat16(self):
+        result = cast_to_bfloat16(self.array_float32)
+        self.assertEqual(result.dtype, jnp.bfloat16)
+        
+        result = cast_to_bfloat16(self.nested_dict)
+        self.assertEqual(result['a'].dtype, jnp.bfloat16)
+        self.assertEqual(result['b']['c'].dtype, jnp.bfloat16)
+        self.assertEqual(result['b']['d'].dtype, jnp.bfloat16)
+
+    def test_cast_to_full_precision(self):
+        result = cast_to_full_precision(self.array_float16)
+        self.assertEqual(result.dtype, jnp.float32)
+        
+        result = cast_to_full_precision(self.nested_dict)
+        self.assertEqual(result['a'].dtype, jnp.float32)
+        self.assertEqual(result['b']['c'].dtype, jnp.float32)
+        self.assertEqual(result['b']['d'].dtype, jnp.float32)
+
+    def test_cast_to_half_precision(self):
+        result = cast_to_half_precision(self.array_float32)
+        self.assertEqual(result.dtype, HALF_PRECISION_DATATYPE)
+        
+        result = cast_to_half_precision(self.nested_dict)
+        self.assertEqual(result['a'].dtype, HALF_PRECISION_DATATYPE)
+        self.assertEqual(result['b']['c'].dtype, HALF_PRECISION_DATATYPE)
+        self.assertEqual(result['b']['d'].dtype, HALF_PRECISION_DATATYPE)
+
+    def test_force_full_precision_decorator(self):
+        @force_full_precision
+        def test_func(x, y):
+            return x + y, x * y
+
+        # Test with float16 inputs
+        x = jnp.array([1.0, 2.0], dtype=jnp.float16)
+        y = jnp.array([3.0, 4.0], dtype=jnp.float16)
+        
+        result1, result2 = test_func(x, y)
+        
+        # Check that inputs were converted to float32 during computation
+        self.assertEqual(result1.dtype, jnp.float16)  # Output is cast back to float16
+        self.assertEqual(result2.dtype, jnp.float16)  # Output is cast back to float16
+
+    def test_mixed_tree_handling(self):
+        # Test that non-array elements are preserved
+        result = cast_to_float32(self.mixed_tree)
+        self.assertEqual(result['array'].dtype, jnp.float32)
+        self.assertEqual(result['scalar'], 42)
+        self.assertEqual(result['nested']['none'], None)
+        self.assertEqual(result['nested']['array'].dtype, jnp.float32)
+
+    def test_empty_structures(self):
+        # Test with empty structures
+        empty_dict = {}
+        result = cast_to_float32(empty_dict)
+        self.assertEqual(result, {})
+
+        empty_list = []
+        result = cast_to_float32(empty_list)
+        self.assertEqual(result, [])
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_dtypes.py

@@ -0,0 +1,20 @@
+import unittest
+import jax.numpy as jnp
+from mpfj.dtypes 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)
+
+    def test_set_half_precision_datatype(self):
+        """Test setting half precision datatype to bfloat16"""
+        set_half_precision_datatype(jnp.bfloat16)
+        self.assertEqual(half_precision_datatype(), jnp.bfloat16)
+        
+        # Reset to default
+        set_half_precision_datatype(jnp.float16)
+        self.assertEqual(half_precision_datatype(), jnp.float16)
+
+if __name__ == '__main__':
+    unittest.main()