Fix JAX CriticMultiplexer type annotation.
"""Multiplexers are networks that take multiple inputs."""
-from typing import Optional
+from typing import Callable, Optional, Union
from acme.jax import utils
import haiku as hk
@@ -23,6 +23,7 @@ import jax.numpy as jnp
import tensorflow_probability
tfd = tensorflow_probability.substrates.jax.distributions
+ModuleOrArrayTransform = Union[hk.Module, Callable[[jnp.ndarray], jnp.ndarray]]
class CriticMultiplexer(hk.Module):
@@ -43,9 +44,9 @@ class CriticMultiplexer(hk.Module):
"""
def __init__(self,
- critic_network: Optional[hk.Transformed] = None,
- observation_network: Optional[hk.Transformed] = None,
- action_network: Optional[hk.Transformed] = None):
+ critic_network: Optional[ModuleOrArrayTransform] = None,
+ observation_network: Optional[ModuleOrArrayTransform] = None,
+ action_network: Optional[ModuleOrArrayTransform] = None):
self._critic_network = critic_network
self._observation_network = observation_network
self._action_network = action_network
