CI: update ruff to v0.6.1

Author: jakevdp

.pre-commit-config.yaml

@@ -26,7 +26,7 @@ repos:
     files: \.py$
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.4.4
+  rev: v0.6.1
   hooks:
   - id: ruff
 

cloud_tpu_colabs/JAX_NeurIPS_2020_demo.ipynb

@@ -510,8 +510,8 @@
    "outputs": [],
    "source": [
     "image_partitions = P(1, 1, 4, 2)\n",
-    "sharded_conv = sharded_jit(conv, \n",
-    "                           in_parts=(image_partitions, None), \n",
+    "sharded_conv = sharded_jit(conv,\n",
+    "                           in_parts=(image_partitions, None),\n",
     "                           out_parts=image_partitions)\n",
     "\n",
     "sharded_conv(image, kernel)"

cloud_tpu_colabs/JAX_demo.ipynb

@@ -877,7 +877,7 @@
     "  def g(z):\n",
     "    return jnp.cos(z) * jnp.tan(y.sum()) * jnp.tanh(x).sum()\n",
     "  return grad(lambda w: jnp.sum(g(w)))(x)\n",
-    "  \n",
+    "\n",
     "f(x)"
    ]
   },
@@ -950,17 +950,6 @@
     "per_example_hess = pmap(input_hess)  # pmap!\n",
     "per_example_hess(inputs)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 0,
-   "metadata": {
-    "colab": {},
-    "colab_type": "code",
-    "id": "u3ggM_WYZ8QC"
-   },
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

cloud_tpu_colabs/Wave_Equation.ipynb

@@ -67,7 +67,6 @@
    "source": [
     "from functools import partial\n",
     "import jax\n",
-    "from jax import jit, pmap\n",
     "from jax import lax\n",
     "from jax import tree_util\n",
     "import jax.numpy as jnp\n",

docs/_tutorials/advanced-autodiff.md

@@ -640,7 +640,7 @@ def our_jacrev(f):
         y, vjp_fun = vjp(f, x)
         # Use vmap to do a matrix-Jacobian product.
         # Here, the matrix is the Euclidean basis, so we get all
-        # entries in the Jacobian at once. 
+        # entries in the Jacobian at once.
         J, = vmap(vjp_fun, in_axes=0)(jnp.eye(len(y)))
         return J
     return jacfun
@@ -654,7 +654,7 @@ from jax import jacfwd as builtin_jacfwd
 def our_jacfwd(f):
     def jacfun(x):
         _jvp = lambda s: jvp(f, (x,), (s,))[1]
-        Jt =vmap(_jvp, in_axes=1)(jnp.eye(len(x)))
+        Jt = vmap(_jvp, in_axes=1)(jnp.eye(len(x)))
         return jnp.transpose(Jt)
     return jacfun
 

docs/jep/9407-type-promotion.ipynb

@@ -3317,7 +3317,6 @@
    ],
    "source": [
     "# @title\n",
-    "from jax import dtypes\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import pandas as pd\n",
@@ -3802,7 +3801,6 @@
    ],
    "source": [
     "# @title\n",
-    "from jax import dtypes\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import pandas as pd\n",

docs/jep/9407-type-promotion.md

@@ -908,7 +908,6 @@ display.HTML(table.to_html())
 :tags: [hide-input]
 
 # @title
-from jax import dtypes
 import jax
 import jax.numpy as jnp
 import pandas as pd
@@ -963,7 +962,6 @@ display.HTML(table.to_html())
 :tags: [hide-input]
 
 # @title
-from jax import dtypes
 import jax
 import jax.numpy as jnp
 import pandas as pd

docs/notebooks/Common_Gotchas_in_JAX.ipynb

@@ -226,7 +226,6 @@
    ],
    "source": [
     "import jax.numpy as jnp\n",
-    "import jax.lax as lax\n",
     "from jax import make_jaxpr\n",
     "\n",
     "# lax.fori_loop\n",
@@ -1031,7 +1030,6 @@
     }
    ],
    "source": [
-    "from jax import random\n",
     "key = random.key(0)\n",
     "key"
    ]
@@ -1105,8 +1103,8 @@
     "print(\"old key\", key)\n",
     "key, subkey = random.split(key)\n",
     "normal_pseudorandom = random.normal(subkey, shape=(1,))\n",
-    "print(\"    \\---SPLIT --> new key   \", key)\n",
-    "print(\"             \\--> new subkey\", subkey, \"--> normal\", normal_pseudorandom)"
+    "print(r\"    \\---SPLIT --> new key   \", key)\n",
+    "print(r\"             \\--> new subkey\", subkey, \"--> normal\", normal_pseudorandom)"
    ]
   },
   {
@@ -1140,8 +1138,8 @@
     "print(\"old key\", key)\n",
     "key, subkey = random.split(key)\n",
     "normal_pseudorandom = random.normal(subkey, shape=(1,))\n",
-    "print(\"    \\---SPLIT --> new key   \", key)\n",
-    "print(\"             \\--> new subkey\", subkey, \"--> normal\", normal_pseudorandom)"
+    "print(r\"    \\---SPLIT --> new key   \", key)\n",
+    "print(r\"             \\--> new subkey\", subkey, \"--> normal\", normal_pseudorandom)"
    ]
   },
   {
@@ -1701,7 +1699,7 @@
    ],
    "source": [
     "init_val = 0\n",
-    "cond_fun = lambda x: x<10\n",
+    "cond_fun = lambda x: x < 10\n",
     "body_fun = lambda x: x+1\n",
     "lax.while_loop(cond_fun, body_fun, init_val)\n",
     "# --> array(10, dtype=int32)"

docs/notebooks/Common_Gotchas_in_JAX.md

@@ -130,7 +130,6 @@ It is not recommended to use iterators in any JAX function you want to `jit` or
 :outputId: 52d885fd-0239-4a08-f5ce-0c38cc008903
 
 import jax.numpy as jnp
-import jax.lax as lax
 from jax import make_jaxpr
 
 # lax.fori_loop
@@ -471,7 +470,6 @@ The random state is described by a special array element that we call a __key__:
 :id: yPHE7KTWgAWs
 :outputId: ae8af0ee-f19e-474e-81b6-45e894eb2fc3
 
-from jax import random
 key = random.key(0)
 key
 ```
@@ -504,8 +502,8 @@ Instead, we __split__ the PRNG to get usable __subkeys__ every time we need a ne
 print("old key", key)
 key, subkey = random.split(key)
 normal_pseudorandom = random.normal(subkey, shape=(1,))
-print("    \---SPLIT --> new key   ", key)
-print("             \--> new subkey", subkey, "--> normal", normal_pseudorandom)
+print(r"    \---SPLIT --> new key   ", key)
+print(r"             \--> new subkey", subkey, "--> normal", normal_pseudorandom)
 ```
 
 +++ {"id": "tqtFVE4MthO3"}
@@ -519,8 +517,8 @@ We propagate the __key__ and make new __subkeys__ whenever we need a new random
 print("old key", key)
 key, subkey = random.split(key)
 normal_pseudorandom = random.normal(subkey, shape=(1,))
-print("    \---SPLIT --> new key   ", key)
-print("             \--> new subkey", subkey, "--> normal", normal_pseudorandom)
+print(r"    \---SPLIT --> new key   ", key)
+print(r"             \--> new subkey", subkey, "--> normal", normal_pseudorandom)
 ```
 
 +++ {"id": "0KLYUluz3lN3"}
@@ -805,7 +803,7 @@ def while_loop(cond_fun, body_fun, init_val):
 :outputId: 552fe42f-4d32-4e25-c8c2-b951160a3f4e
 
 init_val = 0
-cond_fun = lambda x: x<10
+cond_fun = lambda x: x < 10
 body_fun = lambda x: x+1
 lax.while_loop(cond_fun, body_fun, init_val)
 # --> array(10, dtype=int32)

docs/notebooks/Custom_derivative_rules_for_Python_code.ipynb

@@ -247,7 +247,6 @@
     }
    ],
    "source": [
-    "import jax.numpy as jnp\n",
     "\n",
     "def log1pexp(x):\n",
     "  return jnp.log(1. + jnp.exp(x))\n",
@@ -984,7 +983,7 @@
     "                             (a, x_star, x_star_bar),\n",
     "                             x_star_bar))\n",
     "  return a_bar, jnp.zeros_like(x_star)\n",
-    "  \n",
+    "\n",
     "def rev_iter(f, packed, u):\n",
     "  a, x_star, x_star_bar = packed\n",
     "  _, vjp_x = vjp(lambda x: f(a, x), x_star)\n",
@@ -1884,7 +1883,6 @@
     }
    ],
    "source": [
-    "from jax import vjp\n",
     "\n",
     "y, f_vjp = vjp(f, 3.)\n",
     "print(y)"
@@ -1983,7 +1981,7 @@
     "  return x, x\n",
     "\n",
     "def debug_bwd(x, g):\n",
-    "  import pdb; pdb.set_trace()\n",
+    "  pdb.set_trace()\n",
     "  return g\n",
     "\n",
     "debug.defvjp(debug_fwd, debug_bwd)"

docs/notebooks/Custom_derivative_rules_for_Python_code.md

@@ -145,7 +145,6 @@ Say we want to write a function called `log1pexp`, which computes $x \mapsto \lo
 :id: 6lWbTvs40ET-
 :outputId: 8caff99e-add1-4c70-ace3-212c0c5c6f4e
 
-import jax.numpy as jnp
 
 def log1pexp(x):
   return jnp.log(1. + jnp.exp(x))
@@ -524,7 +523,7 @@ def fixed_point_rev(f, res, x_star_bar):
                              (a, x_star, x_star_bar),
                              x_star_bar))
   return a_bar, jnp.zeros_like(x_star)
-  
+
 def rev_iter(f, packed, u):
   a, x_star, x_star_bar = packed
   _, vjp_x = vjp(lambda x: f(a, x), x_star)
@@ -965,7 +964,6 @@ print(grad(f)(3.))
 :id: s1Pn_qCIODcF
 :outputId: 423d34e0-35b8-4b57-e89d-f70f20e28ea9
 
-from jax import vjp
 
 y, f_vjp = vjp(f, 3.)
 print(y)
@@ -1015,7 +1013,7 @@ def debug_fwd(x):
   return x, x
 
 def debug_bwd(x, g):
-  import pdb; pdb.set_trace()
+  pdb.set_trace()
   return g
 
 debug.defvjp(debug_fwd, debug_bwd)

docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb

@@ -30,9 +30,7 @@
    },
    "outputs": [],
    "source": [
-    "import os\n",
     "\n",
-    "import functools\n",
     "from typing import Optional\n",
     "\n",
     "import numpy as np\n",

docs/notebooks/Distributed_arrays_and_automatic_parallelization.md

@@ -26,9 +26,7 @@ This tutorial discusses parallelism via `jax.Array`, the unified array object mo
 ```{code-cell}
 :id: FNxScTfq3vGF
 
-import os
 
-import functools
 from typing import Optional
 
 import numpy as np