nit

ASEM000 · Apr 2, 2024 · 3207225 · 3207225
1 parent 829b431
commit 3207225
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 9 deletions.
diff --git a/serket/_src/nn/convolution.py b/serket/_src/nn/convolution.py
@@ -20,7 +20,7 @@
 import functools as ft
 import operator as op
 from itertools import product
-from typing import Any, Sequence
+from typing import Sequence
 
 import jax
 import jax.numpy as jnp
@@ -42,7 +42,6 @@
     maybe_lazy_call,
     maybe_lazy_init,
     positive_int_cb,
-    single_dispatch,
     validate_in_features_shape,
     validate_spatial_ndim,
 )

diff --git a/serket/_src/nn/linear.py b/serket/_src/nn/linear.py
@@ -28,13 +28,7 @@
     resolve_activation,
 )
 from serket._src.nn.initialization import DType, InitType, resolve_init
-from serket._src.utils import (
-    maybe_lazy_call,
-    maybe_lazy_init,
-    positive_int_cb,
-    single_dispatch,
-    tuplify,
-)
+from serket._src.utils import maybe_lazy_call, maybe_lazy_init, positive_int_cb, tuplify
 
 T = TypeVar("T")
 PyTree = Any