# python3
# Copyright 2018 DeepMind Technologies Limited. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Multiplexers are networks that take multiple inputs."""

from typing import Optional

from acme.jax import utils
import haiku as hk
import jax.numpy as jnp
import tensorflow_probability

tfd = tensorflow_probability.substrates.jax.distributions


class CriticMultiplexer(hk.Module):
  """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()`.
  """

  def __init__(self,
               critic_network: Optional[hk.Transformed] = None,
               observation_network: Optional[hk.Transformed] = None,
               action_network: Optional[hk.Transformed] = None):
    self._critic_network = critic_network
    self._observation_network = observation_network
    self._action_network = action_network
    super().__init__(name='critic_multiplexer')

  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

@@ -697,16 +697,26 @@
 g import
+ Callable,
  Optiona
@@ -716,16 +716,23 @@
 Optional
+, Union
 %0A%0Afrom a
@@ -884,16 +884,96 @@
 butions%0A
+ModuleOrArrayTransform = Union%5Bhk.Module, Callable%5B%5Bjnp.ndarray%5D, jnp.ndarray%5D%5D%0A
 %0A%0Aclass 
@@ -1759,35 +1759,45 @@
 twork: Optional%5B
-hk.
+ModuleOrArray
 Transformed%5D = N
@@ -1781,34 +1781,32 @@
 OrArrayTransform
-ed
 %5D = None,%0A      
@@ -1836,35 +1836,45 @@
 twork: Optional%5B
-hk.
+ModuleOrArray
 Transformed%5D = N
@@ -1858,34 +1858,32 @@
 OrArrayTransform
-ed
 %5D = None,%0A      
@@ -1916,19 +1916,29 @@
 ptional%5B
-hk.
+ModuleOrArray
 Transfor
@@ -1938,18 +1938,16 @@
 ransform
-ed
 %5D = None

Fix JAX CriticMultiplexer type annotation.