From b232d0944071041f739d2015430edbc2ff7c0612 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Thu, 29 Jul 2021 09:51:41 -0400
Subject: [PATCH] Enable flake8 checks for spaces around operators.

---
 jax/_src/image/scale.py     |  2 +-
 jax/_src/lax/lax.py         |  2 +-
 jax/_src/numpy/lax_numpy.py | 10 +++++-----
 jax/_src/scipy/linalg.py    | 14 +++++++-------
 jax/experimental/jet.py     |  6 +++---
 setup.cfg                   |  2 +-
 tests/api_test.py           | 10 +++++-----
 tests/linalg_test.py        |  4 ++--
 tests/random_test.py        |  8 ++++----
 9 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/jax/_src/image/scale.py b/jax/_src/image/scale.py
index b63fd98fc..c1a753bd9 100644
--- a/jax/_src/image/scale.py
+++ b/jax/_src/image/scale.py
@@ -35,7 +35,7 @@ def _fill_keys_cubic_kernel(x):
   # IEEE Transactions on Acoustics, Speech, and Signal Processing,
   # 29(6):1153–1160, 1981.
   out = ((1.5 * x - 2.5) * x) * x + 1.
-  out = jnp.where(x >= 1., ((-0.5* x + 2.5) * x - 4.) * x + 2., out)
+  out = jnp.where(x >= 1., ((-0.5 * x + 2.5) * x - 4.) * x + 2., out)
   return jnp.where(x >= 2., 0., out)
 
 
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index fc0f82f25..17ef566a3 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -1924,7 +1924,7 @@ def collapse(operand: Array, start_dimension: int,
 
 def slice_in_dim(operand: Array, start_index: Optional[int],
                  limit_index: Optional[int],
-                 stride: int = 1, axis: int = 0)-> Array:
+                 stride: int = 1, axis: int = 0) -> Array:
   """Convenience wrapper around slice applying to only one dimension."""
   start_indices = [0] * operand.ndim
   limit_indices = list(operand.shape)
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 43d038c4f..35868b432 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3345,8 +3345,8 @@ def meshgrid(*args, **kwargs):
 
 
 def _make_1d_grid_from_slice(s: slice, op_name: str):
-  start =core.concrete_or_error(None, s.start,
-                                f"slice start of jnp.{op_name}") or 0
+  start = core.concrete_or_error(None, s.start,
+                                 f"slice start of jnp.{op_name}") or 0
   stop = core.concrete_or_error(None, s.stop,
                                 f"slice stop of jnp.{op_name}")
   step = core.concrete_or_error(None, s.step,
@@ -3907,7 +3907,7 @@ def diag(v, k=0):
   else:
     raise ValueError("diag input must be 1d or 2d")
 
-_SCALAR_VALUE_DOC="""\
+_SCALAR_VALUE_DOC = """\
 This differs from np.diagflat for some scalar values of v,
 jax always returns a two-dimensional array, whereas numpy may
 return a scalar depending on the type of v.
@@ -3929,7 +3929,7 @@ def diagflat(v, k=0):
   res = res.reshape(adj_length,adj_length)
   return res
 
-_POLY_DOC="""\
+_POLY_DOC = """\
 This differs from np.poly when an integer array is given.
 np.poly returns a result with dtype float64 in this case.
 jax returns a result with an inexact type, but not necessarily
@@ -4032,7 +4032,7 @@ def trim_zeros(filt, trim='fb'):
   end = argmin(nz[::-1]) if 'b' in trim.lower() else 0
   return filt[start:len(filt) - end]
 
-_LEADING_ZEROS_DOC="""\
+_LEADING_ZEROS_DOC = """\
 Setting trim_leading_zeros=True makes the output match that of numpy.
 But prevents the function from being able to be used in compiled code.
 """
diff --git a/jax/_src/scipy/linalg.py b/jax/_src/scipy/linalg.py
index 65518564d..6d57b274c 100644
--- a/jax/_src/scipy/linalg.py
+++ b/jax/_src/scipy/linalg.py
@@ -288,10 +288,10 @@ def _calc_P_Q(A):
    n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
    A = A / 2**n_squarings
    U13, V13 = _pade13(A)
-   conds=jnp.array([1.495585217958292e-002, 2.539398330063230e-001,
-                    9.504178996162932e-001, 2.097847961257068e+000])
-   U = jnp.select((A_L1<conds), (U3, U5, U7, U9), U13)
-   V = jnp.select((A_L1<conds), (V3, V5, V7, V9), V13)
+   conds = jnp.array([1.495585217958292e-002, 2.539398330063230e-001,
+                      9.504178996162932e-001, 2.097847961257068e+000])
+   U = jnp.select((A_L1 < conds), (U3, U5, U7, U9), U13)
+   V = jnp.select((A_L1 < conds), (V3, V5, V7, V9), V13)
   elif A.dtype == 'float32' or A.dtype == 'complex64':
     U3,V3 = _pade3(A)
     U5,V5 = _pade5(A)
@@ -299,9 +299,9 @@ def _calc_P_Q(A):
     n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
     A = A / 2**n_squarings
     U7,V7 = _pade7(A)
-    conds=jnp.array([4.258730016922831e-001, 1.880152677804762e+000])
-    U = jnp.select((A_L1<conds), (U3, U5), U7)
-    V = jnp.select((A_L1<conds), (V3, V5), V7)
+    conds = jnp.array([4.258730016922831e-001, 1.880152677804762e+000])
+    U = jnp.select((A_L1 < conds), (U3, U5), U7)
+    V = jnp.select((A_L1 < conds), (V3, V5), V7)
   else:
     raise TypeError("A.dtype={} is not supported.".format(A.dtype))
   P = U + V  # p_m(A) : numerator
diff --git a/jax/experimental/jet.py b/jax/experimental/jet.py
index 24f79b15e..c1683b9a0 100644
--- a/jax/experimental/jet.py
+++ b/jax/experimental/jet.py
@@ -364,7 +364,7 @@ def _exp_taylor(primals_in, series_in):
   u = [x] + series
   v = [lax.exp(x)] + [None] * len(series)
   for k in range(1,len(v)):
-    v[k] = fact(k-1) * sum([_scale(k, j)* v[k-j] * u[j] for j in range(1, k+1)])
+    v[k] = fact(k-1) * sum([_scale(k, j) * v[k-j] * u[j] for j in range(1, k+1)])
   primal_out, *series_out = v
   return primal_out, series_out
 jet_rules[lax.exp_p] = _exp_taylor
@@ -377,7 +377,7 @@ def _pow_taylor(primals_in, series_in):
   u = [x] + series
   v = [u_ ** r_] + [None] * len(series)
   for k in range(1, len(v)):
-    v[k] = fact(k-1) * sum([_scale(k, j)* v[k-j] * u[j] for j in range(1, k+1)])
+    v[k] = fact(k-1) * sum([_scale(k, j) * v[k-j] * u[j] for j in range(1, k+1)])
   primal_out, *series_out = v
 
   return primal_out, series_out
@@ -412,7 +412,7 @@ def _expit_taylor(primals_in, series_in):
   e = [v[0] * (1 - v[0])] + [None] * len(series)  # terms for sigmoid' = sigmoid * (1 - sigmoid)
   for k in range(1, len(v)):
     v[k] = fact(k-1) * sum([_scale(k, j) * e[k-j] * u[j] for j in range(1, k+1)])
-    e[k] = (1 - v[0]) * v[k] - fact(k) * sum([_scale2(k, j)* v[j] * v[k-j] for j in range(1, k+1)])
+    e[k] = (1 - v[0]) * v[k] - fact(k) * sum([_scale2(k, j) * v[j] * v[k-j] for j in range(1, k+1)])
 
   primal_out, *series_out = v
   return primal_out, series_out
diff --git a/setup.cfg b/setup.cfg
index 683500d2f..7519e9cd4 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -6,7 +6,7 @@ ignore =
   E121 # line continuations
   W503, W504 # line breaks around binary operators
 max-complexity = 18
-select = B,C,F,W,T4,B9
+select = B,C,F,W,T4,B9,E225,E227,E228
 exclude =
   .git,
   build,
diff --git a/tests/api_test.py b/tests/api_test.py
index dc8d4df23..dbd5a008f 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1222,11 +1222,11 @@ class APITest(jtu.JaxTestCase):
       return jnp.sum(jnp.cos(jnp.abs(z)))
 
     ans = grad(f)(zs)
-    expected = np.array([ 0.        +0.j,
-                          -0.80430663+0.40215331j,
-                          -0.70368982+0.35184491j,
-                           0.1886467 -0.09432335j,
-                           0.86873727-0.43436864j])
+    expected = np.array([ 0.          + 0.j,
+                          -0.80430663 + 0.40215331j,
+                          -0.70368982 + 0.35184491j,
+                           0.1886467  - 0.09432335j,
+                           0.86873727 - 0.43436864j])
     self.assertAllClose(ans, expected, check_dtypes=False,
                         atol=jtu.default_gradient_tolerance,
                         rtol=jtu.default_gradient_tolerance)
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 04149087b..0028ac824 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -440,7 +440,7 @@ class NumpyLinalgTest(jtu.JaxTestCase):
     new_w, new_v = f(new_a)
     new_a = (new_a + np.conj(new_a.T)) / 2
     # Assert rtol eigenvalue delta between perturbed eigenvectors vs new true eigenvalues.
-    RTOL=1e-2
+    RTOL = 1e-2
     assert np.max(
       np.abs((np.diag(np.dot(np.conj((v+dv).T), np.dot(new_a,(v+dv)))) - new_w) / new_w)) < RTOL
     # Redundant to above, but also assert rtol for eigenvector property with new true eigenvalues.
@@ -676,7 +676,7 @@ class NumpyLinalgTest(jtu.JaxTestCase):
 
     # Check that q is close to unitary.
     self.assertTrue(np.all(
-        norm(np.eye(k) -np.matmul(np.conj(T(lq)), lq)) < 5))
+        norm(np.eye(k) - np.matmul(np.conj(T(lq)), lq)) < 5))
 
     if not full_matrices and m >= n:
         jtu.check_jvp(jnp.linalg.qr, partial(jvp, jnp.linalg.qr), (a,), atol=3e-3)
diff --git a/tests/random_test.py b/tests/random_test.py
index 7b031f6fb..78bc03705 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -910,8 +910,8 @@ class LaxRandomTest(jtu.JaxTestCase):
     # 1/2 CDF_one_maxwell((x - loc) / scale))
     #   + 1/2 (1 - CDF_one_maxwell(- (x - loc) / scale)))
     def double_sided_maxwell_cdf(x, loc, scale):
-      pos = scipy.stats.maxwell().cdf((x - loc)/ scale)
-      neg = (1 - scipy.stats.maxwell().cdf((-x + loc)/ scale))
+      pos = scipy.stats.maxwell().cdf((x - loc) / scale)
+      neg = (1 - scipy.stats.maxwell().cdf((-x + loc) / scale))
       return (pos + neg) / 2
 
     for samples in [uncompiled_samples, compiled_samples]:
@@ -939,9 +939,9 @@ class LaxRandomTest(jtu.JaxTestCase):
       assert len(counts) == 2
 
       self.assertAllClose(
-          counts[0]/ num_samples, 0.5, rtol=1e-02, atol=1e-02)
+          counts[0] / num_samples, 0.5, rtol=1e-02, atol=1e-02)
       self.assertAllClose(
-          counts[1]/ num_samples, 0.5, rtol=1e-02, atol=1e-02)
+          counts[1] / num_samples, 0.5, rtol=1e-02, atol=1e-02)
 
   def testChoiceShapeIsNotSequenceError(self):
     key = random.PRNGKey(0)