renamed module to mpx

Author: Alexander

mpfj/__init__.py mpx/__init__.pympfj/__init__.py renamed to mpx/__init__.py

@@ -33,8 +33,8 @@
 
 import optax
 
-import mpfj.cast as cast
-import mpfj.loss_scaling as loss_scaling
+import mpx.cast as cast
+import mpx.loss_scaling as loss_scaling
 
 from jaxtyping import PyTree, Bool
 

mpfj/cast.py mpx/cast.pympfj/cast.py renamed to mpx/cast.py

@@ -5,7 +5,7 @@
 from jaxtyping import Array, Float, Int, PyTree
 import numpy as np
 
-from mpfj.cast import (
+from mpx.cast import (
     cast_tree,
     cast_to_float32,
     cast_to_float16,
@@ -14,7 +14,7 @@
     cast_to_half_precision,
     force_full_precision,
 )
-from mpfj.dtypes import HALF_PRECISION_DATATYPE
+from mpx.dtypes import HALF_PRECISION_DATATYPE
 
 
 class EQXModuleBase(eqx.Module):

mpfj/dtypes.py mpx/dtypes.pympfj/dtypes.py renamed to mpx/dtypes.py

@@ -1,6 +1,6 @@
 import unittest
 import jax.numpy as jnp
-from mpfj.dtypes import half_precision_datatype, set_half_precision_datatype
+from mpx.dtypes import half_precision_datatype, set_half_precision_datatype
 
 class TestDtypes(unittest.TestCase):
     def test_default_half_precision(self):

mpfj/grad_tools.py mpx/grad_tools.pympfj/grad_tools.py renamed to mpx/grad_tools.py

@@ -3,8 +3,8 @@
 import jax.numpy as jnp
 import equinox as eqx
 import optax
-from mpfj.grad_tools import select_tree, filter_grad, filter_value_and_grad, optimizer_update
-from mpfj.loss_scaling import DynamicLossScaling
+from mpx.grad_tools import select_tree, filter_grad, filter_value_and_grad, optimizer_update
+from mpx.loss_scaling import DynamicLossScaling
 
 
 # Create a simple model for testing

mpfj/loss_scaling.py mpx/loss_scaling.pympfj/loss_scaling.py renamed to mpx/loss_scaling.py

@@ -2,7 +2,7 @@
 import jax
 import jax.numpy as jnp
 import equinox as eqx
-from mpfj.loss_scaling import DynamicLossScaling, all_finite, scaled
+from mpx.loss_scaling import DynamicLossScaling, all_finite, scaled
 
 class TestLossScaling(unittest.TestCase):
     def setUp(self):