From 659615ec10e8fb7fe2c258033f7308cfd6eab8ca Mon Sep 17 00:00:00 2001 From: Justin Fu Date: Fri, 18 Oct 2024 15:17:49 -0700 Subject: [PATCH] Resolve linter errors --- jax/_src/api.py | 1 - jax/_src/interpreters/jaxpr_passes.py | 2 +- jax/_src/state/primitives.py | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/jax/_src/api.py b/jax/_src/api.py index 0e3706050c5b..d2ac5465eded 100644 --- a/jax/_src/api.py +++ b/jax/_src/api.py @@ -77,7 +77,6 @@ from jax._src.interpreters import ad from jax._src.interpreters import batching -from jax._src.interpreters import jaxpr_passes from jax._src.interpreters import partial_eval as pe from jax._src.interpreters import pxla from jax._src.interpreters import xla diff --git a/jax/_src/interpreters/jaxpr_passes.py b/jax/_src/interpreters/jaxpr_passes.py index 959aa877ece7..97af5c99fd1f 100644 --- a/jax/_src/interpreters/jaxpr_passes.py +++ b/jax/_src/interpreters/jaxpr_passes.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from collections.abc import Callable, Iterable, Iterator, Sequence +from collections.abc import Callable, Sequence import dataclasses import functools from functools import partial diff --git a/jax/_src/state/primitives.py b/jax/_src/state/primitives.py index ad822f2883c4..57ab6b6b8a81 100644 --- a/jax/_src/state/primitives.py +++ b/jax/_src/state/primitives.py @@ -703,4 +703,4 @@ def _broadcast_to_edtype_rule(ctx: jaxpr_passes.ResolveEdtypesContext, a, *, shape): raise NotImplementedError() -jaxpr_passes.register_edtype_rule(broadcast_to_p, _broadcast_to_edtype_rule) \ No newline at end of file +jaxpr_passes.register_edtype_rule(broadcast_to_p, _broadcast_to_edtype_rule)