Update lock and fix errors

.pre-commit-config.yaml

@@ -15,12 +15,12 @@ repos:
   - id: check-yaml
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.7.4
+  rev: v0.8.2
   hooks:
   - id: ruff
 
 - repo: https://github.com/RobertCraigie/pyright-python
-  rev: v1.1.389
+  rev: v1.1.390
   hooks:
   - id: pyright
 

pyproject.toml

@@ -49,7 +49,7 @@ dev-dependencies = [
   'pylint >= 3.3.1',
   'pyright >= 0.0.13',
   'pytest >= 8',
-  'ruff >= 0.7',
+  'ruff >= 0.8.1',
 ]
 
 [tool.isort]

tjax/_src/cotangent_tools.py

@@ -182,9 +182,9 @@ def _cotangent_combinator_bwd(f: Callable[..., tuple[XT, Y]],
                                                          strict=True)):
         scaled_y_bar = tree.map(lambda y_bar_i, scale=aux_cotangent_scale: y_bar_i * scale,
                                 y_bar)
-        this_xs_bar = cast(XT, (xs_zero[:i]
-                                + (x_bar,)
-                                + xs_zero[i + 1:]))
+        this_xs_bar = cast('XT', (xs_zero[:i]
+                                  + (x_bar,)
+                                  + xs_zero[i + 1:]))
         this_result_bar = (this_xs_bar, scaled_y_bar)
         this_args_bar = f_vjp(this_result_bar)
         all_args_bar.append(this_args_bar)

tjax/_src/dataclasses/dataclass.py

@@ -92,7 +92,7 @@ def f(x: type[Any], /) -> type[TDataclassInstance]:
 
     # Apply dataclass function to cls.
     data_clz: type[TDataclassInstance] = cast(
-            type[TDataclassInstance],
+            'type[TDataclassInstance]',
             dataclasses.dataclass(init=init, repr=repr, eq=eq, order=order, frozen=frozen)(cls))
 
     # Partition fields into static, and dynamic; and assign these to the class.

tjax/_src/graph/register_jax.py

@@ -60,9 +60,9 @@ def flatten_with_keys(graph: nx.Graph[Any], /
             return (zip(keys, values, strict=True), keys)
 
         flatten_with_keys_ = cast(
-                Callable[[Graph], tuple[Iterable[tuple[Hashable, Any]], Hashable]],
+                "Callable[[Graph], tuple[Iterable[tuple[Hashable, Any]], Hashable]]",
                 flatten_with_keys)
-        unflatten_tree_ = cast(Callable[[Hashable, Any], Graph], unflatten_tree)
-        flatten_tree_ = cast(Callable[[Graph], tuple[Iterable[Any], Hashable]],
+        unflatten_tree_ = cast("Callable[[Hashable, Any], Graph]", unflatten_tree)
+        flatten_tree_ = cast("Callable[[Graph], tuple[Iterable[Any], Hashable]]",
                              flatten_tree)
         register_pytree_with_keys(graph_type, flatten_with_keys_, unflatten_tree_, flatten_tree_)

tjax/_src/partial.py

@@ -77,7 +77,7 @@ def tree_unflatten(cls,
         if not isinstance(dynamic_kwargs, dict):
             raise RuntimeError  # noqa: TRY004
 
-        dynamic_kwargs = cast(dict[str, Any], dynamic_kwargs)
+        dynamic_kwargs = cast('dict[str, Any]', dynamic_kwargs)
         args = cls._unpartition_args(static_argnums, static_args, dynamic_args,
                                      callable_is_static=callable_is_static)