Update default package versions used by uv. (#106)

Upgrades jax, jaxlib, flax, and treescope. Also fixes or silences
some pylint warnings that were surfaced by updating the uv
environment.

.pylintrc

@@ -100,6 +100,7 @@ disable=abstract-method,
         long-suffix,
         map-builtin-not-iterating,
         misplaced-comparison-constant,
+        missing-class-docstring,
         missing-function-docstring,
         metaclass-assignment,
         next-method-called,

penzai/core/named_axes.py

@@ -88,11 +88,12 @@ def my_dot(a, b):
 
 import abc
 import collections
+from collections.abc import Callable, Hashable, Mapping, Sequence
 import dataclasses
 import functools
 import operator
 import typing
-from typing import Any, Callable, Hashable, Mapping, Sequence
+from typing import Any
 
 import jax
 import jax.numpy as jnp

penzai/core/struct.py

@@ -23,11 +23,12 @@
 from __future__ import annotations
 
 import abc
+from collections.abc import Callable, Hashable, Sequence
 import dataclasses
 import functools
 import inspect
 import typing
-from typing import Any, Callable, Hashable, Literal, Sequence, Type, TypeVar
+from typing import Any, Literal, Type, TypeVar
 
 import jax
 from typing_extensions import dataclass_transform

penzai/core/variables.py

@@ -69,11 +69,12 @@
 
 import abc
 import contextlib
+from collections.abc import Callable, Hashable, Iterable
 import dataclasses
 import functools
 import inspect
 import typing
-from typing import Any, Callable, Generic, Hashable, Iterable, Literal, TypeAlias, TypeVar
+from typing import Any, Generic, Literal, TypeAlias, TypeVar
 
 import jax
 import jax.numpy as jnp

penzai/nn/attention.py

@@ -23,8 +23,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Hashable
 import dataclasses
-from typing import Any, Hashable
+from typing import Any
 
 import jax
 import jax.numpy as jnp