Add isort pre-commit and run (#68)

python/mlx/extension.py

@@ -6,7 +6,7 @@ import subprocess
 import sys
 from pathlib import Path
 
-from setuptools import Extension, setup, find_namespace_packages
+from setuptools import Extension, find_namespace_packages, setup
 from setuptools.command.build_ext import build_ext
 
 import mlx

python/mlx/nn/__init__.py

@@ -1,5 +1,5 @@
 # Copyright © 2023 Apple Inc.
 
-from mlx.nn.layers import *
 from mlx.nn import losses
+from mlx.nn.layers import *
 from mlx.nn.utils import value_and_grad

python/mlx/nn/layers/__init__.py

@@ -1,6 +1,5 @@
 # Copyright © 2023 Apple Inc.
 
-from mlx.nn.layers.base import Module
 from mlx.nn.layers.activations import (
     GELU,
     ReLU,
@@ -11,6 +10,7 @@ from mlx.nn.layers.activations import (
     relu,
     silu,
 )
+from mlx.nn.layers.base import Module
 from mlx.nn.layers.containers import Sequential
 from mlx.nn.layers.convolution import Conv1d, Conv2d
 from mlx.nn.layers.dropout import Dropout

python/mlx/nn/layers/base.py

@@ -1,7 +1,7 @@
 # Copyright © 2023 Apple Inc.
 
 import textwrap
-from typing import Any, Callable, List, Union, Optional
+from typing import Any, Callable, List, Optional, Union
 
 import mlx.core as mx
 from mlx.utils import tree_flatten, tree_unflatten

python/mlx/nn/layers/transformer.py

@@ -1,7 +1,7 @@
 # Copyright © 2023 Apple Inc.
 
 import math
-from typing import Optional, Any
+from typing import Any, Optional
 
 import mlx.core as mx
 from mlx.nn.layers.base import Module

python/tests/test_array.py

@@ -5,9 +5,8 @@ import unittest
 from itertools import permutations
 
 import mlx.core as mx
-import numpy as np
-
 import mlx_tests
+import numpy as np
 
 
 class TestVersion(mlx_tests.MLXTestCase):

python/tests/test_autograd.py

@@ -3,7 +3,6 @@
 import unittest
 
 import mlx.core as mx
-
 import mlx_tests
 
 

python/tests/test_bf16.py

@@ -1,13 +1,12 @@
 # Copyright © 2023 Apple Inc.
 
+import math
 import unittest
 from itertools import permutations
 
-import math
 import mlx.core as mx
-import numpy as np
-
 import mlx_tests
+import numpy as np
 
 try:
     import torch

python/tests/test_blas.py

@@ -1,13 +1,12 @@
 # Copyright © 2023 Apple Inc.
 
+import math
 import unittest
 from itertools import permutations
 
-import math
 import mlx.core as mx
-import numpy as np
-
 import mlx_tests
+import numpy as np
 
 
 class TestBlas(mlx_tests.MLXTestCase):

python/tests/test_conv.py

@@ -1,13 +1,12 @@
 # Copyright © 2023 Apple Inc.
 
+import math
 import unittest
 from itertools import permutations
 
-import math
 import mlx.core as mx
-import numpy as np
-
 import mlx_tests
+import numpy as np
 
 try:
     import torch

python/tests/test_device.py

@@ -3,7 +3,6 @@
 import unittest
 
 import mlx.core as mx
-
 import mlx_tests
 
 

python/tests/test_eval.py

@@ -1,11 +1,9 @@
 # Copyright © 2023 Apple Inc.
 
+import unittest
 from functools import partial
 
-import unittest
-
 import mlx.core as mx
-
 import mlx_tests
 
 

python/tests/test_fft.py

@@ -1,12 +1,11 @@
 # Copyright © 2023 Apple Inc.
 
+import itertools
 import unittest
 
-import itertools
 import mlx.core as mx
-import numpy as np
-
 import mlx_tests
+import numpy as np
 
 
 class TestFFT(mlx_tests.MLXTestCase):

python/tests/test_load.py

@@ -1,12 +1,12 @@
 # Copyright © 2023 Apple Inc.
 
-import unittest
 import os
-import mlx.core as mx
-import numpy as np
 import tempfile
+import unittest
 
+import mlx.core as mx
 import mlx_tests
+import numpy as np
 
 
 class TestLoad(mlx_tests.MLXTestCase):

python/tests/test_nn.py

@@ -1,15 +1,14 @@
 # Copyright © 2023 Apple Inc.
 
+import os
+import tempfile
 import unittest
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_map, tree_unflatten
-import numpy as np
-import os
-import tempfile
-
 import mlx_tests
+import numpy as np
+from mlx.utils import tree_flatten, tree_map, tree_unflatten
 
 
 class TestNN(mlx_tests.MLXTestCase):

python/tests/test_ops.py

@@ -1,13 +1,12 @@
 # Copyright © 2023 Apple Inc.
 
+import math
 import unittest
 from itertools import permutations
 
-import math
 import mlx.core as mx
-import numpy as np
-
 import mlx_tests
+import numpy as np
 
 
 class TestOps(mlx_tests.MLXTestCase):

python/tests/test_optimizers.py

@@ -5,7 +5,6 @@ import unittest
 import mlx.core as mx
 import mlx.optimizers as opt
 import mlx.utils
-
 import mlx_tests
 
 

python/tests/test_random.py

@@ -3,7 +3,6 @@
 import unittest
 
 import mlx.core as mx
-
 import mlx_tests
 
 

python/tests/test_reduce.py

@@ -1,12 +1,11 @@
 # Copyright © 2023 Apple Inc.
 
 import unittest
-from itertools import permutations, combinations
+from itertools import combinations, permutations
 
 import mlx.core as mx
-import numpy as np
-
 import mlx_tests
+import numpy as np
 
 
 class TestReduce(mlx_tests.MLXTestCase):

python/tests/test_tree.py

@@ -4,7 +4,6 @@ import unittest
 
 import mlx.core as mx
 import mlx.utils
-
 import mlx_tests
 
 

python/tests/test_vmap.py

@@ -3,7 +3,6 @@
 import unittest
 
 import mlx.core as mx
-
 import mlx_tests