Add isort pre-commit and run (#68)

.pre-commit-config.yaml

@@ -7,3 +7,9 @@ repos:
     rev: 22.10.0
     hooks:
     -   id: black
+-   repo: https://github.com/pycqa/isort
+    rev: 5.12.0
+    hooks:
+    -   id: isort
+        args:
+            - --profile=black

benchmarks/numpy/single_ops.py

@@ -1,7 +1,6 @@
 # Copyright © 2023 Apple Inc.
 
 import numpy as np
-
 from time_utils import time_fn
 
 

benchmarks/python/batch_matmul_bench.py

@@ -1,8 +1,8 @@
 # Copyright © 2023 Apple Inc.
 
 import argparse
-import mlx.core as mx
 
+import mlx.core as mx
 from time_utils import time_fn
 
 B = 8

benchmarks/python/blas/bench_gemm.py

@@ -1,13 +1,14 @@
 # Copyright © 2023 Apple Inc.
 
-import numpy as np
 import argparse
-import mlx.core as mx
-import time
-import torch
-import os
 import math
+import os
 import subprocess
+import time
+
+import mlx.core as mx
+import numpy as np
+import torch
 
 device_name = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
 device_name = device_name.decode("utf-8").strip("\n")

benchmarks/python/blas/bench_gemv.py

@@ -1,14 +1,14 @@
 # Copyright © 2023 Apple Inc.
 
-import matplotlib.pyplot as plt
-import numpy as np
 import argparse
-import mlx.core as mx
-import time
-import torch
 import os
 import subprocess
+import time
 
+import matplotlib.pyplot as plt
+import mlx.core as mx
+import numpy as np
+import torch
 
 results_dir = "./results"
 

benchmarks/python/llama_torch_bench.py

@@ -4,8 +4,8 @@ import math
 import time
 
 import torch
-import torch.nn as nn
 import torch.mps
+import torch.nn as nn
 
 
 def sync_if_needed(x):

benchmarks/python/single_ops.py

@@ -1,8 +1,8 @@
 # Copyright © 2023 Apple Inc.
 
 import argparse
-import mlx.core as mx
 
+import mlx.core as mx
 from time_utils import time_fn
 
 

examples/extensions/mlx_sample_extensions/__init__.py

@@ -1,4 +1,5 @@
 # Copyright © 2023 Apple Inc.
 
 import mlx.core as mx
+
 from .mlx_sample_extensions import *

examples/extensions/setup.py

@@ -1,8 +1,9 @@
 # Copyright © 2023 Apple Inc.
 
-from mlx import extension
 from setuptools import setup
 
+from mlx import extension
+
 if __name__ == "__main__":
     setup(
         name="mlx_sample_extensions",

examples/python/linear_regression.py

@@ -1,8 +1,9 @@
 # Copyright © 2023 Apple Inc.
 
-import mlx.core as mx
 import time
 
+import mlx.core as mx
+
 num_features = 100
 num_examples = 1_000
 num_iters = 10_000

examples/python/logistic_regression.py

@@ -1,8 +1,9 @@
 # Copyright © 2023 Apple Inc.
 
-import mlx.core as mx
 import time
 
+import mlx.core as mx
+
 num_features = 100
 num_examples = 1_000
 num_iters = 10_000

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
 
 

setup.py

@@ -9,7 +9,7 @@ import sysconfig
 from pathlib import Path
 from subprocess import run
 
-from setuptools import Extension, setup, find_namespace_packages
+from setuptools import Extension, find_namespace_packages, setup
 from setuptools.command.build_ext import build_ext