From 3207225e1bc4c765b1383ac7b93049ed14d31b11 Mon Sep 17 00:00:00 2001 From: ASEM000 Date: Wed, 3 Apr 2024 01:07:00 +0900 Subject: [PATCH] nit --- serket/_src/nn/convolution.py | 3 +-- serket/_src/nn/linear.py | 8 +------- 2 files changed, 2 insertions(+), 9 deletions(-) diff --git a/serket/_src/nn/convolution.py b/serket/_src/nn/convolution.py index 41355c1..cf5d37a 100644 --- 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 index 8926ae6..a13136c 100644 --- 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