class_docstr: Module connecting a critic torso to (transformed) observations/actions.
  
  This takes as input a `critic_network`, an `observation_network`, and an
  `action_network` and returns another network whose outputs are given by
  `critic_network(observation_network(o), action_network(a))`.
  
  The observations and actions passed to this module are assumed to have a batch
  dimension that match.
  
  Notes:
  - Either the `observation_` or `action_network` can be `None`, in which case
    the observation or action, resp., are passed to the critic network as is.
  - If all `critic_`, `observation_` and `action_network` are `None`, this
    module reduces to a simple `tf2_utils.batch_concat()`.
code: def __call__(
      self, observation: jnp.ndarray, action: jnp.ndarray
  ) -> jnp.ndarray:
      # Maybe transform observations and actions before feeding them on.
      if self._observation_network:
          observation = self._observation_network(observation)
      if self._action_network:
          action = self._action_network(action)
  
      # Concat observations and actions, with one batch dimension.
      outputs = utils.batch_concat([observation, action])
  
      # Maybe transform output before returning.
      if self._critic_network:
          outputs = self._critic_network(outputs)
  
      return outputs
  
file_path: acme/acme/jax/networks/multiplexers.py
imports: from typing import Callable, Optional, Union
  from acme.jax import utils
  import haiku
  import jax.numpy
  import tensorflow_probability
label_desc: Use image node_id 2 for calling the CriticMultiplexer obj's underlying member method code with example usage: obj.__call__(observation, action) and returns: outputs