Resolve linter errors

jax-ml · Oct 18, 2024 · 659615e · 659615e
1 parent dbe5cae
commit 659615e
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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)
+jaxpr_passes.register_edtype_rule(broadcast_to_p, _broadcast_to_edtype_rule)