precommit

graphcore-research · Nov 8, 2023 · 62b2289 · 62b2289
1 parent 9611bb3
commit 62b2289
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/jax_scaled_arithmetics/interpreters/autoscale.py b/jax_scaled_arithmetics/interpreters/autoscale.py
@@ -1,12 +1,13 @@
 # Copyright (c) 2023 Graphcore Ltd. All rights reserved.
 
+from functools import wraps
+
 import jax
 import numpy as np
-
 from jax import core
 from jax._src.util import safe_map
+
 from ..lax import scaled_ops_registry
-from functools import wraps
 
 
 def autoscale(fun):

diff --git a/jax_scaled_arithmetics/lax/__init__.py b/jax_scaled_arithmetics/lax/__init__.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2023 Graphcore Ltd. All rights reserved.
 
 from jax import lax
+
 from .scaled_ops import *
 
 scaled_ops_registry = {lax.mul_p: scaled_mul}