diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py
index 5a80ab4c..ba513d96 100644
--- a/mjx/mujoco/mjx/_src/dataclasses.py
+++ b/mjx/mujoco/mjx/_src/dataclasses.py
@@ -18,7 +18,7 @@ import copy
 import dataclasses
 
 import typing
-from typing import Any, Dict, Optional, Sequence, TypeVar, Union
+from typing import Dict, Optional, Sequence, Tuple, TypeVar, Union
 import jax
 import numpy as np
 
@@ -122,7 +122,7 @@ class PyTreeNode:
     raise NotImplementedError
 
   @classmethod
-  def fields(cls) -> tuple[dataclasses.Field[Any], ...]:
+  def fields(cls) -> Tuple[dataclasses.Field, ...]:  # pylint: disable=g-bare-generic
     return dataclasses.fields(cls)
 
   def tree_replace(
diff --git a/mjx/mujoco/mjx/_src/test_util.py b/mjx/mujoco/mjx/_src/test_util.py
index 36d07d9f..a02d3922 100644
--- a/mjx/mujoco/mjx/_src/test_util.py
+++ b/mjx/mujoco/mjx/_src/test_util.py
@@ -17,7 +17,7 @@
 import os
 import sys
 import time
-from typing import Dict, Tuple
+from typing import Dict, Optional, Tuple
 from xml.etree import ElementTree as ET
 
 from etils import epath
@@ -201,9 +201,9 @@ def _make_geom(
 
 def _make_actuator(
     actuator_type: str,
-    joint: str | None = None,
-    site: str | None = None,
-    refsite: str | None = None,
+    joint: Optional[str] = None,
+    site: Optional[str] = None,
+    refsite: Optional[str] = None,
 ) -> Dict[str, str]:
   """Returns attributes for an actuator."""
   if joint: