style(isort): equinox and jaxtyping are first-party

docs/examples/custom_rules.ipynb

@@ -22,12 +22,12 @@
     "import functools as ft\n",
     "from typing import Union\n",
     "\n",
-    "import equinox as eqx  # /~https://github.com/patrick-kidger/equinox\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
-    "from jaxtyping import ArrayLike  # /~https://github.com/patrick-kidger/jaxtyping\n",
     "\n",
-    "import quax"
+    "import equinox as eqx  # /~https://github.com/patrick-kidger/equinox\n",
+    "import quax\n",
+    "from jaxtyping import ArrayLike  # /~https://github.com/patrick-kidger/jaxtyping"
    ]
   },
   {

docs/examples/default_rules.ipynb

@@ -26,14 +26,14 @@
     "from collections.abc import Sequence\n",
     "from typing import Union\n",
     "\n",
-    "import equinox as eqx  # /~https://github.com/patrick-kidger/equinox\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import jax.random as jr\n",
     "import jax.tree_util as jtu\n",
-    "from jaxtyping import ArrayLike  # /~https://github.com/patrick-kidger/quax\n",
     "\n",
-    "import quax"
+    "import equinox as eqx  # /~https://github.com/patrick-kidger/equinox\n",
+    "import quax\n",
+    "from jaxtyping import ArrayLike  # /~https://github.com/patrick-kidger/quax"
    ]
   },
   {

docs/examples/redispatch.ipynb

@@ -21,17 +21,17 @@
    "source": [
     "from typing import Union\n",
     "\n",
-    "import equinox as eqx  # /~https://github.com/patrick-kidger/equinox\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
+    "\n",
+    "import equinox as eqx  # /~https://github.com/patrick-kidger/equinox\n",
+    "import quax\n",
     "from jaxtyping import (  # /~https://github.com/patrick-kidger/quax\n",
     "    Array,\n",
     "    ArrayLike,\n",
     "    Int,\n",
     "    Shaped,\n",
-    ")\n",
-    "\n",
-    "import quax"
+    ")"
    ]
   },
   {

pyproject.toml

@@ -30,6 +30,12 @@ dependencies = [
     "plum-dispatch>=2.2.1",
 ]
 
+[dependency-groups]
+dev = [
+    "ipykernel>=6.29.5",
+    "pre-commit>=4.1.0",
+]
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
@@ -38,6 +44,7 @@ build-backend = "hatchling.build"
 include = ["quax/*"]
 
 [tool.pytest.ini_options]
+minversion = 8.0
 addopts = "--jaxtyping-packages=quax,beartype.beartype(conf=beartype.BeartypeConf(strategy=beartype.BeartypeStrategy.On))"
 
 [tool.pytest_env]
@@ -55,6 +62,7 @@ fixable = ["I001", "F401"]
 combine-as-imports = true
 lines-after-imports = 2
 extra-standard-library = ["typing_extensions"]
+known-first-party = ["equinox", "jaxtyping"]
 order-by-type = false
 
 [tool.pyright]

quax/_core.py

@@ -5,7 +5,6 @@
 from typing import Any, cast, Generic, TypeVar, Union
 from typing_extensions import TypeGuard
 
-import equinox as eqx
 import jax
 import jax._src
 import jax.core as core
@@ -15,6 +14,8 @@
 import jax.tree_util as jtu
 import plum
 from jax.custom_derivatives import SymbolicZero as SZ
+
+import equinox as eqx
 from jaxtyping import ArrayLike, PyTree
 
 

quax/examples/lora/_core.py

@@ -1,14 +1,14 @@
 from typing import Union
 
-import equinox as eqx
 import jax.core
 import jax.lax as lax
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu
-from jaxtyping import Array, ArrayLike, PRNGKeyArray, PyTree, Shaped
 
+import equinox as eqx
 import quax
+from jaxtyping import Array, ArrayLike, PRNGKeyArray, PyTree, Shaped
 
 
 class LoraArray(quax.ArrayValue):

quax/examples/named/_core.py

@@ -2,14 +2,14 @@
 from collections.abc import Callable
 from typing import Any, Generic, Optional, TypeVar, Union
 
-import equinox as eqx
 import jax.core
 import jax.extend as jex
 import jax.lax as lax
 import jax.numpy as jnp
-from jaxtyping import ArrayLike
 
+import equinox as eqx
 import quax
+from jaxtyping import ArrayLike
 
 
 @dataclasses.dataclass(frozen=True, eq=False)

quax/examples/prng/_core.py

@@ -4,17 +4,17 @@
 from typing import Any, TypeVar
 from typing_extensions import Self, TYPE_CHECKING, TypeAlias
 
-import equinox as eqx
 import jax
 import jax._src.prng
 import jax.core
 import jax.lax as lax
 import jax.numpy as jnp
 import jax.tree_util as jtu
 import numpy as np
-from jaxtyping import Array, ArrayLike, Float, Integer, UInt, UInt32
 
+import equinox as eqx
 import quax
+from jaxtyping import Array, ArrayLike, Float, Integer, UInt, UInt32
 
 
 RealArray: TypeAlias = ArrayLike

quax/examples/sparse/_core.py

@@ -1,12 +1,12 @@
 from typing import get_args
 
-import equinox as eqx
 import jax.core
 import jax.lax as lax
 import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike, Integer, Shaped
 
+import equinox as eqx
 import quax
+from jaxtyping import Array, ArrayLike, Integer, Shaped
 
 
 class BCOO(quax.ArrayValue):

quax/examples/structured_matrices/_core.py

@@ -1,12 +1,12 @@
 from typing import Union
 
-import equinox as eqx
 import jax.core
 import jax.lax as lax
 import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike, Shaped
 
+import equinox as eqx
 import quax
+from jaxtyping import Array, ArrayLike, Shaped
 
 
 class TridiagonalMatrix(quax.ArrayValue):

quax/examples/unitful/_core.py

@@ -1,12 +1,12 @@
 from typing import Union
 
-import equinox as eqx  # /~https://github.com/patrick-kidger/equinox
 import jax
 import jax.core as core
 import jax.numpy as jnp
-from jaxtyping import ArrayLike  # /~https://github.com/patrick-kidger/jaxtyping
 
+import equinox as eqx  # /~https://github.com/patrick-kidger/equinox
 import quax
+from jaxtyping import ArrayLike  # /~https://github.com/patrick-kidger/jaxtyping
 
 
 class Dimension:

quax/examples/zero/_core.py

@@ -1,13 +1,13 @@
 import functools as ft
 from typing import Any, get_args, Union
 
-import equinox as eqx
 import jax.core
 import jax.lax as lax
 import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike
 
+import equinox as eqx
 import quax
+from jaxtyping import Array, ArrayLike
 
 
 class Zero(quax.ArrayValue):

tests/conftest.py

@@ -1,7 +1,7 @@
-import equinox.internal as eqxi
+import jax.random as jr
 import pytest
 
 
 @pytest.fixture()
 def getkey():
-    return eqxi.GetKey()
+    return lambda: jr.PRNGKey(0)

tests/test_core.py

@@ -1,14 +1,14 @@
 from typing import cast
 
-import equinox as eqx
 import jax
 import jax.core
 import jax.lax as lax
 import jax.numpy as jnp
 import pytest
-from jaxtyping import Array
 
+import equinox as eqx
 import quax
+from jaxtyping import Array
 
 
 def test_jit_inline():

tests/test_lora.py

@@ -1,14 +1,14 @@
-import equinox as eqx
 import jax
 import jax.lax as lax
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
-from jaxtyping import TypeCheckError
 from plum import NotFoundLookupError
 
+import equinox as eqx
 import quax
 import quax.examples.lora as lora
+from jaxtyping import TypeCheckError
 
 
 def test_linear(getkey):

tests/test_named.py

@@ -1,14 +1,14 @@
 from typing import cast
 
-import equinox as eqx
 import jax.lax as lax
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
-from jaxtyping import Array
 
+import equinox as eqx
 import quax
 import quax.examples.named as named
+from jaxtyping import Array
 
 
 def test_init(getkey):

tests/test_sparse.py

@@ -1,8 +1,8 @@
-import equinox as eqx
 import jax
 import jax.numpy as jnp
 import jax.random as jr
 
+import equinox as eqx
 import quax
 import quax.examples.sparse as sparse
 

tests/test_zero.py

@@ -1,10 +1,10 @@
-import equinox as eqx
 import jax
 import jax.lax as lax
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
 
+import equinox as eqx
 import quax
 import quax.examples.zero as zero