From c6ab354bd7e7ac6c963f521f6092a9b919ed4086 Mon Sep 17 00:00:00 2001
From: Erik Frey <erikfrey@google.com>
Date: Mon, 22 Apr 2024 16:21:44 -0700
Subject: [PATCH] Remove type annotations unavailable in python < 3.10.  Fixes
 #1590

PiperOrigin-RevId: 627188913
Change-Id: I9adc627349a4001291ac6d05080b761fee155ab0
---
 mjx/mujoco/mjx/_src/dataclasses.py | 4 ++--
 mjx/mujoco/mjx/_src/test_util.py   | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

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: