Remove type annotations unavailable in python < 3.10. Fixes #1590

PiperOrigin-RevId: 627188913
Change-Id: I9adc627349a4001291ac6d05080b761fee155ab0

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(

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,
+    joint: Optional[str] = None,
-    site: str | None = None,
+    site: Optional[str] = None,
-    refsite: str | None = None,
+    refsite: Optional[str] = None,
 ) -> Dict[str, str]:
   """Returns attributes for an actuator."""
   if joint: