diff --git a/sequence_layers/mlx/__init__.py b/sequence_layers/mlx/__init__.py
index 5587b4b..8cb7355 100644
--- a/sequence_layers/mlx/__init__.py
+++ b/sequence_layers/mlx/__init__.py
@@ -13,4 +13,4 @@
 # limitations under the License.
 """Sequence layers in MLX."""
 
-from sequence_layers.mlx.basic_types import *
\ No newline at end of file
+from sequence_layers.mlx.types import *
\ No newline at end of file
diff --git a/sequence_layers/mlx/basic_types.py b/sequence_layers/mlx/types.py
similarity index 100%
rename from sequence_layers/mlx/basic_types.py
rename to sequence_layers/mlx/types.py
diff --git a/sequence_layers/mlx/types_test.py b/sequence_layers/mlx/types_test.py
index a9279aa..60be1d7 100644
--- a/sequence_layers/mlx/types_test.py
+++ b/sequence_layers/mlx/types_test.py
@@ -1,6 +1,6 @@
 import mlx.core as mx
 import numpy as np
-import sequence_layers.mlx.basic_types as types
+from sequence_layers.mlx import types
 from absl.testing import parameterized
 from absl.testing import absltest