diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 0695eff..d97d19f 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,5 +1,10 @@
 exclude: '^ethosu/vela/(tflite|ethos_u55_regs)/'
 repos:
+- repo: https://github.com/asottile/reorder_python_imports
+  rev: v2.2.0
+  hooks:
+  - id: reorder-python-imports
+
 - repo: https://github.com/ambv/black
   rev: stable
   hooks:
diff --git a/README.md b/README.md
index 4d7db1a..4c779d5 100644
--- a/README.md
+++ b/README.md
@@ -110,6 +110,7 @@
 
 ## License
 Vela is licensed under [Apache License 2.0](LICENSE.txt)
+
 ## Contributions and Pull Requests
 
 Contributions are accepted under Apache-2.0. Only submit contributions where you have authored all of the code.
@@ -117,7 +118,7 @@
 ### Sanity checks
 
 The Python codebase is PEP8 compliant with the exception of 120 characters line length.
-We run black and flake8 against the code base excluding "ethosu/vela/tflite/" and "ethosu/vela/ethos\_u55\_regs" directories because they are auto-generated by third party tools.
+We run reorder-python-import, black and flake8 against the code base excluding "ethosu/vela/tflite/" and "ethosu/vela/ethos\_u55\_regs" directories because they are auto-generated by third party tools.
 Those tools are run using [pre-commit framework](https://pre-commit.com/). The configuration file is .pre-commit-config.yaml
 
 #### Install tools
@@ -146,9 +147,17 @@
 
 Those checks can be run manually. This can be achievied running the following
 ```
+$ pre-commit run reorder-python-imports --all-files
+...
 $ pre-commit run flake8 --all-files
 ...
 $ pre-commit run black --all-files
 ```
 
 If you don't specify anything after run, it will execute all the checks.
+```
+$ pre-commit run --all-files
+Reorder python imports...................................................Passed
+black....................................................................Passed
+flake8...................................................................Passed
+```
diff --git a/ethosu/mlw_codec/test_mlw_codec.py b/ethosu/mlw_codec/test_mlw_codec.py
index b868721..0bcd417 100644
--- a/ethosu/mlw_codec/test_mlw_codec.py
+++ b/ethosu/mlw_codec/test_mlw_codec.py
@@ -1,5 +1,4 @@
 #!/usr/bin/env python3
-
 # Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
@@ -15,15 +14,11 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Simple example of the usage of mlw_codec.
-
 import sys
 
 from ethosu import mlw_codec
 
-
 # Simple example
 if __name__ == "__main__":
     weights = [0, 2, 3, 0, -1, -2, -3, 0, 0, 0, 1, -250, 240] * 3
diff --git a/ethosu/vela/__init__.py b/ethosu/vela/__init__.py
index 07d8d79..90376be 100644
--- a/ethosu/vela/__init__.py
+++ b/ethosu/vela/__init__.py
@@ -13,7 +13,6 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 from ._version import __version__
 from .vela import main
 
diff --git a/ethosu/vela/__main__.py b/ethosu/vela/__main__.py
index 9bf74c7..4e35cfc 100644
--- a/ethosu/vela/__main__.py
+++ b/ethosu/vela/__main__.py
@@ -13,7 +13,6 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import sys
 
 from .vela import main
diff --git a/ethosu/vela/_version.py b/ethosu/vela/_version.py
index b670819..8a15dbd 100644
--- a/ethosu/vela/_version.py
+++ b/ethosu/vela/_version.py
@@ -13,7 +13,6 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import pkg_resources
 
 __version__ = pkg_resources.get_distribution("ethos-u-vela").version
diff --git a/ethosu/vela/architecture_features.py b/ethosu/vela/architecture_features.py
index 69f95fa..c8827db 100644
--- a/ethosu/vela/architecture_features.py
+++ b/ethosu/vela/architecture_features.py
@@ -13,22 +13,21 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Holds a container for Ethos-U55/System architecture parameters.
-
 import enum
 from collections import namedtuple
 from configparser import ConfigParser
 
 import numpy as np
 
-from .tensor import MemArea, TensorPurpose, TensorFormat
+from .numeric_util import round_up
+from .numeric_util import round_up_divide
 from .operation import NpuBlockType
-from .numeric_util import round_up, round_up_divide
 from .supported_operators import SupportedOperators
-
+from .tensor import MemArea
+from .tensor import TensorFormat
+from .tensor import TensorPurpose
 
 PointXY = namedtuple("PointXY", "x y")
 PointXYZ = namedtuple("PointXYZ", "x y z")
diff --git a/ethosu/vela/compiler_driver.py b/ethosu/vela/compiler_driver.py
index 6fc3b65..64aff06 100644
--- a/ethosu/vela/compiler_driver.py
+++ b/ethosu/vela/compiler_driver.py
@@ -13,29 +13,27 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Contains the main sequencing of the compiler.
-
 import time
 
+from . import extract_npu_subgraphs
 from . import graph_optimiser
-from . import mark_tensors
+from . import high_level_command_stream_generator
 from . import insert_dma
+from . import live_range
+from . import mark_tensors
+from . import npu_performance
+from . import npu_serialisation
 from . import pass_packing
+from . import register_command_stream_generator
 from . import scheduler
 from . import tensor_allocation
-from . import npu_performance
-from . import high_level_command_stream_generator
-from . import register_command_stream_generator
-from . import extract_npu_subgraphs
-from . import npu_serialisation
 from . import weight_compressor
-from . import live_range
-from .tensor import MemArea
-from .nn_graph import TensorAllocator, PassPlacement
+from .nn_graph import PassPlacement
+from .nn_graph import TensorAllocator
 from .rewrite_graph import verify_graph_health
+from .tensor import MemArea
 
 
 class CompilerOptions:
diff --git a/ethosu/vela/data_type.py b/ethosu/vela/data_type.py
index 6dfe216..bb4c558 100644
--- a/ethosu/vela/data_type.py
+++ b/ethosu/vela/data_type.py
@@ -13,11 +13,8 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Defines the basic numeric type classes for tensors.
-
 import enum
 
 from .numeric_util import round_up_divide
diff --git a/ethosu/vela/driver_actions.py b/ethosu/vela/driver_actions.py
index bd15af2..79ac11a 100644
--- a/ethosu/vela/driver_actions.py
+++ b/ethosu/vela/driver_actions.py
@@ -13,16 +13,15 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Creates driver actions that are embedded in the custom operator payload.
-
 from typing import List
 
 import numpy as np
 
-from .ethos_u55_regs.ethos_u55_regs import config_r, id_r, ARCH_VER
+from .ethos_u55_regs.ethos_u55_regs import ARCH_VER
+from .ethos_u55_regs.ethos_u55_regs import config_r
+from .ethos_u55_regs.ethos_u55_regs import id_r
 
 
 class DACommands:
diff --git a/ethosu/vela/extract_npu_subgraphs.py b/ethosu/vela/extract_npu_subgraphs.py
index ab3db21..6747ec9 100644
--- a/ethosu/vela/extract_npu_subgraphs.py
+++ b/ethosu/vela/extract_npu_subgraphs.py
@@ -13,8 +13,6 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Vela separates CPU operations and NPU operations into separate internal subgraphs. The CPU operations are left
 # untouched in the final output.
@@ -22,11 +20,13 @@
 # Vela does this by identifying NPU passes and pulling them out from the main CPU graph into separate subgraphs, invoked
 # by NpuOp operations. Later, Vela generates command streams and compressed weight streams for the NPU subgraphs and
 # attaches them to the NpuOp. This encapsulates everything the NPU subgraph is supposed to do.
-
 import numpy as np
 
-from .nn_graph import Pass, PassPlacement, Subgraph
-from .operation import Operation, NpuBlockType
+from .nn_graph import Pass
+from .nn_graph import PassPlacement
+from .nn_graph import Subgraph
+from .operation import NpuBlockType
+from .operation import Operation
 
 
 def make_npu_call_op_pass(npu_subgraph):
diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index fdd6fc6..b2b233e 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -13,21 +13,18 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Early optimisation of the network graph, using the rewrite_graph module to do the traversal of the graph. These are
 # split into two parts optimise_graph_a and optimise_graph_b.
-
 import math
 
 import numpy as np
 
 from . import rewrite_graph
-from .operation import Operation, NpuBlockType
-from .tensor import Tensor
 from .data_type import DataType
-
+from .operation import NpuBlockType
+from .operation import Operation
+from .tensor import Tensor
 
 passthrough_nodes = set(("Identity",))
 
diff --git a/ethosu/vela/greedy_allocation.py b/ethosu/vela/greedy_allocation.py
index 6b3d2c1..d6896a5 100644
--- a/ethosu/vela/greedy_allocation.py
+++ b/ethosu/vela/greedy_allocation.py
@@ -13,11 +13,8 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Allocate tensor addresses using a greedy algorithm.
-
 from . import numeric_util
 
 
diff --git a/ethosu/vela/high_level_command_stream.py b/ethosu/vela/high_level_command_stream.py
index bdb0490..2c77e10 100644
--- a/ethosu/vela/high_level_command_stream.py
+++ b/ethosu/vela/high_level_command_stream.py
@@ -13,18 +13,16 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Contains classes that hold commands for the high-level command stream (one command per DMA or NPU stripe).
-
 from enum import IntEnum
 
 import numpy as np
 
-from .operation import NpuBlockType
 from .numeric_util import round_up_divide
-from .range_set import MemoryAccessSet, AccessDirection
+from .operation import NpuBlockType
+from .range_set import AccessDirection
+from .range_set import MemoryAccessSet
 
 
 class Box:
diff --git a/ethosu/vela/high_level_command_stream_generator.py b/ethosu/vela/high_level_command_stream_generator.py
index 47392c0..ef21e06 100644
--- a/ethosu/vela/high_level_command_stream_generator.py
+++ b/ethosu/vela/high_level_command_stream_generator.py
@@ -13,17 +13,17 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Generate a high-level command stream from a scheduled subgraph with CascadedPasses.
 #
 # Also used during scheduling to work out allowable IFM/OFM overlap, this functionality can be accessed using
 # calc_allowed_ofm_ifm_overlap_for_cascaded_pass().
-
-from .nn_graph import SchedulingStrategy, PassPlacement
+from .high_level_command_stream import Box
+from .high_level_command_stream import DMA
+from .high_level_command_stream import NpuStripe
+from .nn_graph import PassPlacement
+from .nn_graph import SchedulingStrategy
 from .operation import NpuBlockType
-from .high_level_command_stream import Box, NpuStripe, DMA
 
 
 def need_dma(tens):
diff --git a/ethosu/vela/insert_dma.py b/ethosu/vela/insert_dma.py
index 33f1a02..703ab9d 100644
--- a/ethosu/vela/insert_dma.py
+++ b/ethosu/vela/insert_dma.py
@@ -13,14 +13,13 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Insert DMA operations into the graph for transfering weights.
-
 from . import rewrite_graph
-from .tensor import MemArea, TensorPurpose
-from .operation import Operation, NpuBlockType
+from .operation import NpuBlockType
+from .operation import Operation
+from .tensor import MemArea
+from .tensor import TensorPurpose
 
 
 def insert_dma_cmd(op, arch):
diff --git a/ethosu/vela/live_range.py b/ethosu/vela/live_range.py
index 54c15ba..23ab67d 100644
--- a/ethosu/vela/live_range.py
+++ b/ethosu/vela/live_range.py
@@ -13,15 +13,13 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Build a live range graph for tensors in one or more subgraphs. Used for tensor allocation as well as in the scheduler.
 # Can work with either a pass packed subgraph or a scheduled subgraph.
-
-from .tensor import Tensor, MemArea
-from .nn_graph import PassPlacement
 from .high_level_command_stream_generator import calc_allowed_ofm_ifm_overlap_for_cascaded_pass
+from .nn_graph import PassPlacement
+from .tensor import MemArea
+from .tensor import Tensor
 
 
 class LiveRange:
diff --git a/ethosu/vela/mark_tensors.py b/ethosu/vela/mark_tensors.py
index c42a28d..3650876 100644
--- a/ethosu/vela/mark_tensors.py
+++ b/ethosu/vela/mark_tensors.py
@@ -13,16 +13,14 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Mark purpose and select formats for Tensors. Also compresses the weights.
-
 from . import rewrite_graph
 from . import weight_compressor
 from .architecture_features import Block
-from .tensor import TensorPurpose, TensorFormat
 from .operation import NpuBlockType
+from .tensor import TensorFormat
+from .tensor import TensorPurpose
 
 
 def purpose_from_list(lst):
diff --git a/ethosu/vela/model_reader.py b/ethosu/vela/model_reader.py
index 6d7a3a4..d1cdc9b 100644
--- a/ethosu/vela/model_reader.py
+++ b/ethosu/vela/model_reader.py
@@ -13,8 +13,6 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Dispatcher for reading a neural network model.
 
diff --git a/ethosu/vela/nn_graph.py b/ethosu/vela/nn_graph.py
index e7820fe..ed2ab32 100644
--- a/ethosu/vela/nn_graph.py
+++ b/ethosu/vela/nn_graph.py
@@ -13,8 +13,6 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Neural network graph classes and enums.
 # Pass - A packed pass containing one or more Operations.
@@ -22,7 +20,6 @@
 #                configurations.
 # Subgraph - Holds a neural network subgraph, pointing at Tensors, Operations, Passes, and CascadedPasses.
 # Graph - A full neural network graph with one or more Subgraphs.
-
 import enum
 
 
diff --git a/ethosu/vela/npu_performance.py b/ethosu/vela/npu_performance.py
index 11f1e92..32208c9 100644
--- a/ethosu/vela/npu_performance.py
+++ b/ethosu/vela/npu_performance.py
@@ -13,24 +13,26 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # NPU performance estimation functions to estimate performance of a Pass and CascadedPass. Uses a model that takes the
 # maximum of the 'cycles required for bandwidth' and 'cycles required for computing'.
 #
 # Called during scheduling to evaluate different proposals, as well as post-scheduling to provide a final performance
 # estimate.
-
 import enum
 
 import numpy as np
 
 from . import numeric_util
-from .tensor import TensorPurpose, MemArea, shape_num_elements, TensorBlockTraversal
-from .nn_graph import PassPlacement, SchedulerRewrite
+from .architecture_features import Block
+from .architecture_features import Kernel
+from .nn_graph import PassPlacement
+from .nn_graph import SchedulerRewrite
 from .operation import NpuBlockType
-from .architecture_features import Block, Kernel
+from .tensor import MemArea
+from .tensor import shape_num_elements
+from .tensor import TensorBlockTraversal
+from .tensor import TensorPurpose
 
 
 def rolling_buffer_dims_from_passes(arch, ps1, block_config_ps1, ps2, block_config_ps2):
diff --git a/ethosu/vela/npu_serialisation.py b/ethosu/vela/npu_serialisation.py
index 29ede84..b8ac20f 100644
--- a/ethosu/vela/npu_serialisation.py
+++ b/ethosu/vela/npu_serialisation.py
@@ -13,20 +13,20 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Serialises and packs an NPU subgraph into tensors.
-
 import struct
 
 import numpy as np
 
 from . import driver_actions
-from .nn_graph import PassPlacement
-from .tensor import MemArea, Tensor, TensorPurpose, TensorFormat
-from .operation import Operation
 from .data_type import DataType
+from .nn_graph import PassPlacement
+from .operation import Operation
+from .tensor import MemArea
+from .tensor import Tensor
+from .tensor import TensorFormat
+from .tensor import TensorPurpose
 
 
 def make_memory_tensor(name, mem_area, sz, want_values, arch):
diff --git a/ethosu/vela/numeric_util.py b/ethosu/vela/numeric_util.py
index 4e61b4c..1722adc 100644
--- a/ethosu/vela/numeric_util.py
+++ b/ethosu/vela/numeric_util.py
@@ -13,11 +13,8 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Numerical utilities for various types of rounding etc.
-
 import math
 
 import numpy as np
diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py
index a25574a..3c776dc 100644
--- a/ethosu/vela/operation.py
+++ b/ethosu/vela/operation.py
@@ -13,11 +13,8 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Internal representation of a Neural Network Operation.
-
 import enum
 
 
diff --git a/ethosu/vela/pass_packing.py b/ethosu/vela/pass_packing.py
index 1ad5b4f..9bf080e 100644
--- a/ethosu/vela/pass_packing.py
+++ b/ethosu/vela/pass_packing.py
@@ -13,17 +13,16 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Packs a subgraph with Neural Network Operations into Passes. Each Pass has one or more Operations.
-
-import enum
 import collections
+import enum
 
-from .nn_graph import Pass, PassPlacement
+from .nn_graph import Pass
+from .nn_graph import PassPlacement
+from .operation import NpuBlockType
+from .operation import Operation
 from .tensor import TensorPurpose
-from .operation import Operation, NpuBlockType
 
 
 class PassFlags(enum.Flag):
diff --git a/ethosu/vela/range_set.py b/ethosu/vela/range_set.py
index d7623c5..f03174e 100644
--- a/ethosu/vela/range_set.py
+++ b/ethosu/vela/range_set.py
@@ -13,11 +13,8 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Helper classes to track memory accesses for calculating dependencies between Commands.
-
 from enum import IntEnum
 from functools import lru_cache
 
diff --git a/ethosu/vela/register_command_stream_generator.py b/ethosu/vela/register_command_stream_generator.py
index 7a4faa8..1489860 100644
--- a/ethosu/vela/register_command_stream_generator.py
+++ b/ethosu/vela/register_command_stream_generator.py
@@ -13,28 +13,44 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Register level (low-level) command stream generation for Ethos-U55. Takes a high-level command stream and generates
 # all the register settings. Calculates dependencies between commands and inserts wait operations. And generates a bit
 # stream suitable for interpretation by the Ethos-U55 processor.
-
 from collections import defaultdict
-from enum import Enum, IntEnum
+from enum import Enum
+from enum import IntEnum
 
 import numpy as np
 
 from . import scaling
+from .architecture_features import ArchitectureFeatures
+from .architecture_features import Block
+from .architecture_features import Kernel
+from .architecture_features import Rect
+from .architecture_features import SharedBufferArea
+from .architecture_features import SHRAMElements
+from .data_type import BaseType
+from .data_type import DataType
+from .ethos_u55_regs.ethos_u55_regs import acc_format
+from .ethos_u55_regs.ethos_u55_regs import activation
+from .ethos_u55_regs.ethos_u55_regs import cmd0
+from .ethos_u55_regs.ethos_u55_regs import cmd1
+from .ethos_u55_regs.ethos_u55_regs import elementwise_mode
+from .ethos_u55_regs.ethos_u55_regs import ifm_precision
+from .ethos_u55_regs.ethos_u55_regs import rounding
 from .high_level_command_stream import CommandType
-from .ethos_u55_regs.ethos_u55_regs import cmd0, cmd1, acc_format, elementwise_mode, rounding, activation, ifm_precision
-from .tensor import MemArea, TensorBlockTraversal, TensorFormat
+from .numeric_util import clamp_sigmoid
+from .numeric_util import clamp_tanh
+from .numeric_util import quantise_float32
+from .numeric_util import round_away_zero
+from .numeric_util import round_up
+from .numeric_util import round_up_to_int
 from .operation import NpuBlockType
-from .numeric_util import quantise_float32, round_up, round_away_zero, round_up_to_int, clamp_sigmoid, clamp_tanh
-from .data_type import BaseType, DataType
 from .shared_buffer_allocation import SharedBufferAllocation
-from .architecture_features import SharedBufferArea, SHRAMElements, ArchitectureFeatures
-from .architecture_features import Block, Kernel, Rect
+from .tensor import MemArea
+from .tensor import TensorBlockTraversal
+from .tensor import TensorFormat
 
 
 class RegisterMachine:
diff --git a/ethosu/vela/rewrite_graph.py b/ethosu/vela/rewrite_graph.py
index e6e24e6..4f0d010 100644
--- a/ethosu/vela/rewrite_graph.py
+++ b/ethosu/vela/rewrite_graph.py
@@ -13,8 +13,6 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Functions for abstracting out the traversal and rewriting of graphs so that the optimisation passes can focus on the
 # correct operation.
diff --git a/ethosu/vela/scaling.py b/ethosu/vela/scaling.py
index 3b749dd..7019c21 100644
--- a/ethosu/vela/scaling.py
+++ b/ethosu/vela/scaling.py
@@ -13,11 +13,8 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Contains various scaling calculations for weights, elementwise operations, pooling etc.
-
 import math
 from enum import IntEnum
 
diff --git a/ethosu/vela/scheduler.py b/ethosu/vela/scheduler.py
index fe31a46..ca018d2 100644
--- a/ethosu/vela/scheduler.py
+++ b/ethosu/vela/scheduler.py
@@ -13,29 +13,34 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # The scheduler costs various strategies for scheduling the network in order to select the block configuration.
-
-import enum
 import copy
+import enum
+from functools import lru_cache
 
 import numpy as np
 
 from . import live_range
 from . import npu_performance
 from . import stats_writer
-from .tensor import TensorPurpose, TensorSubPurpose, TensorFormat, MemArea
-from .operation import NpuBlockType
-from .nn_graph import SchedulingStrategy, CascadedPass, PassPlacement, SchedulerRewrite
-from .npu_performance import make_bandwidth_array, make_macs_array, make_cycles_array, make_metrics_arrays, PassCycles
 from .high_level_command_stream_generator import calc_allowed_ofm_ifm_overlap_for_pass_list
-from .shared_buffer_allocation import (
-    find_block_configs_suitable_for_pass_and_shared_buffer,
-    shared_buffer_allocation_for_pass_and_block_config,
-)
-from functools import lru_cache
+from .nn_graph import CascadedPass
+from .nn_graph import PassPlacement
+from .nn_graph import SchedulerRewrite
+from .nn_graph import SchedulingStrategy
+from .npu_performance import make_bandwidth_array
+from .npu_performance import make_cycles_array
+from .npu_performance import make_macs_array
+from .npu_performance import make_metrics_arrays
+from .npu_performance import PassCycles
+from .operation import NpuBlockType
+from .shared_buffer_allocation import find_block_configs_suitable_for_pass_and_shared_buffer
+from .shared_buffer_allocation import shared_buffer_allocation_for_pass_and_block_config
+from .tensor import MemArea
+from .tensor import TensorFormat
+from .tensor import TensorPurpose
+from .tensor import TensorSubPurpose
 
 
 class ParetoMetric(enum.Enum):
diff --git a/ethosu/vela/shared_buffer_allocation.py b/ethosu/vela/shared_buffer_allocation.py
index 29be6d8..335b863 100644
--- a/ethosu/vela/shared_buffer_allocation.py
+++ b/ethosu/vela/shared_buffer_allocation.py
@@ -13,15 +13,16 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Shared buffer allocation works out how to allocate the Ethos-U55 shared buffer for a given pass.
-
 import numpy as np
 
+from .architecture_features import ArchitectureFeatures
+from .architecture_features import Block
+from .architecture_features import Kernel
+from .architecture_features import SharedBufferArea
+from .architecture_features import SHRAMElements
 from .operation import NpuBlockType
-from .architecture_features import Block, Kernel, SHRAMElements, SharedBufferArea, ArchitectureFeatures
 
 
 class SharedBufferAllocation:
diff --git a/ethosu/vela/stats_writer.py b/ethosu/vela/stats_writer.py
index 3fd29d1..9bbb9db 100644
--- a/ethosu/vela/stats_writer.py
+++ b/ethosu/vela/stats_writer.py
@@ -13,20 +13,20 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Writes out per-pass and summary performance statistics to CSV files.
-
 import csv
 import sys
 
 import numpy as np
 
-from .tensor import MemArea, TensorPurpose
 from .nn_graph import PassPlacement
-from .npu_performance import PassCycles, MacCount, BandwidthDirection
+from .npu_performance import BandwidthDirection
+from .npu_performance import MacCount
+from .npu_performance import PassCycles
 from .numeric_util import round_up_to_int
+from .tensor import MemArea
+from .tensor import TensorPurpose
 
 
 def write_summary_metrics_csv(nng, summary_filename, arch):
diff --git a/ethosu/vela/supported_operators.py b/ethosu/vela/supported_operators.py
index 7334fe2..70700e7 100644
--- a/ethosu/vela/supported_operators.py
+++ b/ethosu/vela/supported_operators.py
@@ -13,11 +13,8 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # The SupportedOperators class which is a collection of all supported operators and parameter checks.
-
 from .data_type import BaseType
 
 
diff --git a/ethosu/vela/tensor.py b/ethosu/vela/tensor.py
index 5cebf4d..19258b5 100644
--- a/ethosu/vela/tensor.py
+++ b/ethosu/vela/tensor.py
@@ -13,19 +13,16 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Internal representation of a Neural Network Tensor.
-
 import enum
 import uuid
 
 import numpy as np
 
 from . import numeric_util
-from .range_set import MemoryRangeSet
 from .numeric_util import round_up_divide
+from .range_set import MemoryRangeSet
 
 
 class MemArea(enum.IntFlag):
diff --git a/ethosu/vela/tensor_allocation.py b/ethosu/vela/tensor_allocation.py
index 255156e..cd2b570 100644
--- a/ethosu/vela/tensor_allocation.py
+++ b/ethosu/vela/tensor_allocation.py
@@ -13,21 +13,18 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Wrapping function to do tensor address allocation. That is, assigning addresses to tensors based on what has been
 # worked out from the allowable overlaps that are calculated by the live range analysis.
-
 import math
 
 import numpy as np
 
 from . import live_range
 from . import numeric_util
-from .tensor import MemArea
-from .nn_graph import TensorAllocator
 from .greedy_allocation import allocate_live_ranges as greedy_allocate_live_ranges
+from .nn_graph import TensorAllocator
+from .tensor import MemArea
 
 
 def linear_allocate_live_ranges(live_ranges, alloc_granularity=256):
diff --git a/ethosu/vela/tflite_mapping.py b/ethosu/vela/tflite_mapping.py
index e8b40bd..4873ecc 100644
--- a/ethosu/vela/tflite_mapping.py
+++ b/ethosu/vela/tflite_mapping.py
@@ -13,122 +13,119 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # TensorFlow Lite mapping functions used by both reader and writer.
 # Contains a mapping from the various TensorFlow Lite enums and options structs, generated by the FlatBuffer code
 # generator, to Vela's internal format.
-
 import struct
 
 import numpy as np
 
 from .data_type import DataType
-from .tflite import Conv2DOptions
-from .tflite import DepthwiseConv2DOptions
-from .tflite import ConcatEmbeddingsOptions
-from .tflite import LSHProjectionOptions
-from .tflite import Pool2DOptions
-from .tflite import SVDFOptions
-from .tflite import RNNOptions
-from .tflite import FullyConnectedOptions
-from .tflite import SoftmaxOptions
-from .tflite import ConcatenationOptions
+from .tflite import AbsOptions
+from .tflite import AddNOptions
 from .tflite import AddOptions
-from .tflite import L2NormOptions
-from .tflite import LocalResponseNormalizationOptions
-from .tflite import LSTMOptions
-from .tflite import ResizeBilinearOptions
-from .tflite import CallOptions
-from .tflite import ReshapeOptions
-from .tflite import SkipGramOptions
-from .tflite import SpaceToDepthOptions
-from .tflite import EmbeddingLookupSparseOptions
-from .tflite import MulOptions
-from .tflite import PadOptions
-from .tflite import GatherOptions
-from .tflite import BatchToSpaceNDOptions
-from .tflite import SpaceToBatchNDOptions
-from .tflite import TransposeOptions
-from .tflite import ReducerOptions
-from .tflite import SubOptions
-from .tflite import DivOptions
-from .tflite import SqueezeOptions
-from .tflite import SequenceRNNOptions
-from .tflite import StridedSliceOptions
-from .tflite import ExpOptions
-from .tflite import TopKV2Options
-from .tflite import SplitOptions
-from .tflite import LogSoftmaxOptions
-from .tflite import CastOptions
-from .tflite import DequantizeOptions
-from .tflite import MaximumMinimumOptions
 from .tflite import ArgMaxOptions
-from .tflite import LessOptions
-from .tflite import NegOptions
-from .tflite import PadV2Options
-from .tflite import GreaterOptions
-from .tflite import GreaterEqualOptions
-from .tflite import LessEqualOptions
-from .tflite import SelectOptions
-from .tflite import SliceOptions
-from .tflite import TransposeConvOptions
-from .tflite import SparseToDenseOptions
-from .tflite import TileOptions
-from .tflite import ExpandDimsOptions
-from .tflite import EqualOptions
-from .tflite import NotEqualOptions
-from .tflite import ShapeOptions
-from .tflite import PowOptions
 from .tflite import ArgMinOptions
-from .tflite import FakeQuantOptions
-from .tflite import PackOptions
-from .tflite import LogicalOrOptions
-from .tflite import OneHotOptions
-from .tflite import LogicalAndOptions
-from .tflite import LogicalNotOptions
-from .tflite import UnpackOptions
-from .tflite import FloorDivOptions
-from .tflite import SquareOptions
-from .tflite import ZerosLikeOptions
-from .tflite import FillOptions
+from .tflite import BatchToSpaceNDOptions
 from .tflite import BidirectionalSequenceLSTMOptions
 from .tflite import BidirectionalSequenceRNNOptions
-from .tflite import UnidirectionalSequenceLSTMOptions
-from .tflite import FloorModOptions
-from .tflite import RangeOptions
-from .tflite import ResizeNearestNeighborOptions
-from .tflite import LeakyReluOptions
-from .tflite import SquaredDifferenceOptions
-from .tflite import MirrorPadOptions
-from .tflite import AbsOptions
-from .tflite import SplitVOptions
-from .tflite import UniqueOptions
-from .tflite import ReverseV2Options
-from .tflite import AddNOptions
-from .tflite import GatherNdOptions
+from .tflite import CallOptions
+from .tflite import CastOptions
+from .tflite import ConcatEmbeddingsOptions
+from .tflite import ConcatenationOptions
+from .tflite import Conv2DOptions
 from .tflite import CosOptions
-from .tflite import WhereOptions
-from .tflite import RankOptions
-from .tflite import ReverseSequenceOptions
-from .tflite import MatrixDiagOptions
-from .tflite import QuantizeOptions
-from .tflite import MatrixSetDiagOptions
 from .tflite import DensifyOptions
 from .tflite import DepthToSpaceOptions
+from .tflite import DepthwiseConv2DOptions
+from .tflite import DequantizeOptions
+from .tflite import DivOptions
+from .tflite import EmbeddingLookupSparseOptions
+from .tflite import EqualOptions
+from .tflite import ExpandDimsOptions
+from .tflite import ExpOptions
+from .tflite import FakeQuantOptions
+from .tflite import FillOptions
+from .tflite import FloorDivOptions
+from .tflite import FloorModOptions
+from .tflite import FullyConnectedOptions
+from .tflite import GatherNdOptions
+from .tflite import GatherOptions
+from .tflite import GreaterEqualOptions
+from .tflite import GreaterOptions
 from .tflite import IfOptions
+from .tflite import L2NormOptions
+from .tflite import LeakyReluOptions
+from .tflite import LessEqualOptions
+from .tflite import LessOptions
+from .tflite import LocalResponseNormalizationOptions
+from .tflite import LogicalAndOptions
+from .tflite import LogicalNotOptions
+from .tflite import LogicalOrOptions
+from .tflite import LogSoftmaxOptions
+from .tflite import LSHProjectionOptions
+from .tflite import LSTMOptions
+from .tflite import MatrixDiagOptions
+from .tflite import MatrixSetDiagOptions
+from .tflite import MaximumMinimumOptions
+from .tflite import MirrorPadOptions
+from .tflite import MulOptions
+from .tflite import NegOptions
 from .tflite import NonMaxSuppressionV4Options
 from .tflite import NonMaxSuppressionV5Options
+from .tflite import NotEqualOptions
+from .tflite import OneHotOptions
+from .tflite import PackOptions
+from .tflite import PadOptions
+from .tflite import PadV2Options
+from .tflite import Pool2DOptions
+from .tflite import PowOptions
+from .tflite import QuantizeOptions
+from .tflite import RangeOptions
+from .tflite import RankOptions
+from .tflite import ReducerOptions
+from .tflite import ReshapeOptions
+from .tflite import ResizeBilinearOptions
+from .tflite import ResizeNearestNeighborOptions
+from .tflite import ReverseSequenceOptions
+from .tflite import ReverseV2Options
+from .tflite import RNNOptions
 from .tflite import ScatterNdOptions
 from .tflite import SegmentSumOptions
+from .tflite import SelectOptions
 from .tflite import SelectV2Options
+from .tflite import SequenceRNNOptions
+from .tflite import ShapeOptions
+from .tflite import SkipGramOptions
+from .tflite import SliceOptions
+from .tflite import SoftmaxOptions
+from .tflite import SpaceToBatchNDOptions
+from .tflite import SpaceToDepthOptions
+from .tflite import SparseToDenseOptions
+from .tflite import SplitOptions
+from .tflite import SplitVOptions
+from .tflite import SquaredDifferenceOptions
+from .tflite import SquareOptions
+from .tflite import SqueezeOptions
+from .tflite import StridedSliceOptions
+from .tflite import SubOptions
+from .tflite import SVDFOptions
+from .tflite import TileOptions
+from .tflite import TopKV2Options
+from .tflite import TransposeConvOptions
+from .tflite import TransposeOptions
+from .tflite import UnidirectionalSequenceLSTMOptions
+from .tflite import UniqueOptions
+from .tflite import UnpackOptions
+from .tflite import WhereOptions
 from .tflite import WhileOptions
-from .tflite.TensorType import TensorType
+from .tflite import ZerosLikeOptions
+from .tflite.ActivationFunctionType import ActivationFunctionType
 from .tflite.BuiltinOperator import BuiltinOperator
 from .tflite.BuiltinOptions import BuiltinOptions
 from .tflite.Padding import Padding
-from .tflite.ActivationFunctionType import ActivationFunctionType
+from .tflite.TensorType import TensorType
 
 
 def inverse_map(map):
diff --git a/ethosu/vela/tflite_reader.py b/ethosu/vela/tflite_reader.py
index aa0ec4d..3632f82 100644
--- a/ethosu/vela/tflite_reader.py
+++ b/ethosu/vela/tflite_reader.py
@@ -13,21 +13,23 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Functions used to read from a TensorFlow Lite format file.
-
 import os.path
 
 import numpy as np
 
-from .tflite.Model import Model
-from .tflite.BuiltinOperator import BuiltinOperator
-from .nn_graph import Graph, Subgraph
+from .nn_graph import Graph
+from .nn_graph import Subgraph
 from .operation import Operation
-from .tensor import Tensor, QuantizationParameters
-from .tflite_mapping import builtin_operator_map, datatype_map, datatype_map_numpy, DataType
+from .tensor import QuantizationParameters
+from .tensor import Tensor
+from .tflite.BuiltinOperator import BuiltinOperator
+from .tflite.Model import Model
+from .tflite_mapping import builtin_operator_map
+from .tflite_mapping import DataType
+from .tflite_mapping import datatype_map
+from .tflite_mapping import datatype_map_numpy
 
 
 def decode_str(s):
diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py
index 1f07242..99df849 100644
--- a/ethosu/vela/tflite_writer.py
+++ b/ethosu/vela/tflite_writer.py
@@ -13,30 +13,31 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Functions used to write to a TensorFlow Lite format file. Supports adding in file identifiers.
-
-import numpy as np
 import flatbuffers
+import flatbuffers.number_types as N
+import numpy as np
+from flatbuffers import encode
 from flatbuffers.builder import UOffsetTFlags
 
-# ugh, the python flatbuffer interface is missing a method to add in file identifier. patching it in here:
-import flatbuffers.number_types as N
-from flatbuffers import encode
-
-from .tflite import Tensor
-from .tflite import QuantizationParameters
-from .tflite import Model
-from .tflite import SubGraph
-from .tflite import OperatorCode
-from .tflite import Operator
+from .nn_graph import PassPlacement
+from .tensor import MemArea
+from .tensor import TensorPurpose
 from .tflite import Buffer
 from .tflite import Metadata
-from .tflite_mapping import datatype_inv_map, builtin_operator_inv_map, custom_prefix, BuiltinOperator
-from .nn_graph import PassPlacement
-from .tensor import TensorPurpose, MemArea
+from .tflite import Model
+from .tflite import Operator
+from .tflite import OperatorCode
+from .tflite import QuantizationParameters
+from .tflite import SubGraph
+from .tflite import Tensor
+from .tflite_mapping import builtin_operator_inv_map
+from .tflite_mapping import BuiltinOperator
+from .tflite_mapping import custom_prefix
+from .tflite_mapping import datatype_inv_map
+
+# ugh, the python flatbuffer interface is missing a method to add in file identifier. patching it in here:
 
 tflite_version = 3
 tflite_file_identifier = "TFL" + str(tflite_version)
diff --git a/ethosu/vela/vela.py b/ethosu/vela/vela.py
index 07772e6..49f8c26 100644
--- a/ethosu/vela/vela.py
+++ b/ethosu/vela/vela.py
@@ -13,30 +13,27 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Main entry point for the Vela compiler.
 #
 # Provides command line interface, options parsing, and network loading. Before calling the compiler driver.
-
-import os
+import argparse
+import ast
+import configparser
 import os.path
 import sys
 import time
-import configparser
-import argparse
-import ast
 
 from . import architecture_features
+from . import compiler_driver
+from . import model_reader
+from . import scheduler
 from . import stats_writer
 from . import tflite_writer
-from . import model_reader
-from . import compiler_driver
-from . import scheduler
 from ._version import __version__
+from .nn_graph import PassPlacement
+from .nn_graph import TensorAllocator
 from .scheduler import ParetoMetric
-from .nn_graph import TensorAllocator, PassPlacement
 from .tensor import MemArea
 
 
diff --git a/ethosu/vela/weight_compressor.py b/ethosu/vela/weight_compressor.py
index ee554b5..04d684e 100644
--- a/ethosu/vela/weight_compressor.py
+++ b/ethosu/vela/weight_compressor.py
@@ -13,26 +13,26 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Compresses and pads the weigths. It also calculates the scales and packs with the biases.
-
 import math
 from collections import namedtuple
 
 import numpy as np
-
-from .numeric_util import round_up
-from .scaling import quantise_scale, reduced_quantise_scale
-from .tensor import TensorPurpose, TensorSubPurpose, TensorFormat, TensorBlockTraversal
-from .operation import NpuBlockType
-from .architecture_features import Block
-from .nn_graph import SchedulingStrategy
-from .data_type import DataType
-
 from ethosu import mlw_codec
 
+from .architecture_features import Block
+from .data_type import DataType
+from .nn_graph import SchedulingStrategy
+from .numeric_util import round_up
+from .operation import NpuBlockType
+from .scaling import quantise_scale
+from .scaling import reduced_quantise_scale
+from .tensor import TensorBlockTraversal
+from .tensor import TensorFormat
+from .tensor import TensorPurpose
+from .tensor import TensorSubPurpose
+
 
 def encode(weight_stream):
     assert np.amin(weight_stream) >= -255
diff --git a/setup.py b/setup.py
index 1a1ae84..8d88997 100644
--- a/setup.py
+++ b/setup.py
@@ -13,13 +13,13 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 # Description:
 # Packaging for the Vela compiler
-
 from os import path
-from setuptools import setup, find_namespace_packages, Extension
+
+from setuptools import Extension
+from setuptools import find_namespace_packages
+from setuptools import setup
 
 # Read the contents of README.md file
 this_directory = path.abspath(path.dirname(__file__))