Make tuple index constexpr-safe; fix Python Protocol import; tighten XLA typings

envpool/core/tuple_utils.h

@@ -28,12 +28,24 @@ struct Index;
 
 template <class T, class... Types>
 struct Index<T, std::tuple<Types...>> {
-  static_assert((std::is_same_v<T, Types> || ...), "Type not found in tuple");
-  static constexpr std::size_t kValue = [] {
-    std::size_t i = 0;
-    ((std::is_same_v<T, Types> ? false : (++i, true)) && ...);
-    return i;
-  }();
+ private:
+  template <std::size_t I>
+  static constexpr std::size_t FindIndex() {
+    return I;
+  }
+
+  template <std::size_t I, class Head, class... Tail>
+  static constexpr std::size_t FindIndex() {
+    if constexpr (std::is_same_v<T, Head>) {
+      return I;
+    } else {
+      return FindIndex<I + 1, Tail...>();
+    }
+  }
+
+ public:
+  static constexpr std::size_t kValue = FindIndex<0, Types...>();
+  static_assert(kValue < sizeof...(Types), "Type not found in tuple");
 };
 
 template <class F, class K, class V, std::size_t... I>

envpool/python/protocol.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Protocol of C++ EnvPool."""
 
+import sys
 from typing import (
   Any,
   Callable,
@@ -30,10 +31,10 @@
 import numpy as np
 from dm_env import TimeStep
 
-try:
+if sys.version_info >= (3, 8):
   from typing import Protocol
-except ImportError:
-  from typing_extensions import Protocol  # type: ignore
+else:
+  from typing_extensions import Protocol
 
 
 class EnvSpec(Protocol):

envpool/python/xla_template.py

@@ -45,8 +45,8 @@ def _normalize_specs(
 
 
 def _make_xla_function(
-  obj: Any, handle: bytes, name: str, specs: Tuple[Tuple[Any], Tuple[Any]],
-  capsules: Tuple[Any, Any]
+  obj: Any, handle: bytes, name: str,
+  specs: Tuple[Tuple[Any, ...], Tuple[Any, ...]], capsules: Tuple[Any, Any]
 ) -> Callable:
   in_specs, out_specs = specs
   in_specs = _normalize_specs(in_specs)