text: # Copyright 2021 The TensorFlow Authors. 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.
  
  """Contains common building blocks for neural networks."""
  from typing import Callable, Dict, List, Optional, Tuple, Union
  
  from absl import logging
  import tensorflow as tf
  
  from official.modeling import tf_utils
  
  
  # Type annotations.
  States = Dict[str, tf.Tensor]
  Activation = Union[str, Callable]
  
  
  def make_divisible(value: float,
                     divisor: int,
                     min_value: Optional[float] = None
                     ) -> int:
    """This is to ensure that all layers have channels that are divisible by 8.
  
    Args:
      value: A `float` of original value.
      divisor: An `int` of the divisor that need to be checked upon.
      min_value: A `float` of  minimum value threshold.
  
    Returns:
      The adjusted value in `int` that is divisible against divisor.
    """
    if min_value is None:
      min_value = divisor
    new_value = max(min_value, int(value + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_value < 0.9 * value:
      new_value += divisor
    return int(new_value)
  
  
  def round_filters(filters: int,
                    multiplier: float,
                    divisor: int = 8,
                    min_depth: Optional[int] = None,
                    skip: bool = False):
    """Rounds number of filters based on width multiplier."""
    orig_f = filters
    if skip or not multiplier:
      return filters
  
    new_filters = make_divisible(value=filters * multiplier,
                                 divisor=divisor,
                                 min_value=min_depth)
  
    logging.info('round_filter input=%s output=%s', orig_f, new_filters)
    return int(new_filters)
  
  
  def hard_swish(x: tf.Tensor) -> tf.Tensor:
    """A Swish6/H-Swish activation function.
  
    Reference: Section 5.2 of Howard et al. "Searching for MobileNet V3."
    https://arxiv.org/pdf/1905.02244.pdf
  
    Args:
      x: the input tensor.
  
    Returns:
      The activation output.
    """
    return x * tf.nn.relu6(x + 3.) * (1. / 6.)
  
  tf.keras.utils.get_custom_objects().update({'hard_swish': hard_swish})
  
  
  def simple_swish(x: tf.Tensor) -> tf.Tensor:
    """A swish/silu activation function without custom gradients.
  
    Useful for exporting to SavedModel to avoid custom gradient warnings.
  
    Args:
      x: the input tensor.
  
    Returns:
      The activation output.
    """
    return x * tf.math.sigmoid(x)
  
  tf.keras.utils.get_custom_objects().update({'simple_swish': simple_swish})
  
  
  @tf.keras.utils.register_keras_serializable(package='Vision')
  class SqueezeExcitation(tf.keras.layers.Layer):
    """Creates a squeeze and excitation layer."""
  
    def __init__(self,
                 in_filters,
                 out_filters,
                 se_ratio,
                 divisible_by=1,
                 use_3d_input=False,
                 kernel_initializer='VarianceScaling',
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activation='relu',
                 gating_activation='sigmoid',
                 **kwargs):
      """Initializes a squeeze and excitation layer.
  
      Args:
        in_filters: An `int` number of filters of the input tensor.
        out_filters: An `int` number of filters of the output tensor.
        se_ratio: A `float` or None. If not None, se ratio for the squeeze and
          excitation layer.
        divisible_by: An `int` that ensures all inner dimensions are divisible by
          this number.
        use_3d_input: A `bool` of whether input is 2D or 3D image.
        kernel_initializer: A `str` of kernel_initializer for convolutional
          layers.
        kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
          Conv2D. Default to None.
        bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2d.
          Default to None.
        activation: A `str` name of the activation function.
        gating_activation: A `str` name of the activation function for final
          gating function.
        **kwargs: Additional keyword arguments to be passed.
      """
      super(SqueezeExcitation, self).__init__(**kwargs)
  
      self._in_filters = in_filters
      self._out_filters = out_filters
      self._se_ratio = se_ratio
      self._divisible_by = divisible_by
      self._use_3d_input = use_3d_input
      self._activation = activation
      self._gating_activation = gating_activation
      self._kernel_initializer = kernel_initializer
      self._kernel_regularizer = kernel_regularizer
      self._bias_regularizer = bias_regularizer
      if tf.keras.backend.image_data_format() == 'channels_last':
        if not use_3d_input:
          self._spatial_axis = [1, 2]
        else:
          self._spatial_axis = [1, 2, 3]
      else:
        if not use_3d_input:
          self._spatial_axis = [2, 3]
        else:
          self._spatial_axis = [2, 3, 4]
      self._activation_fn = tf_utils.get_activation(activation)
      self._gating_activation_fn = tf_utils.get_activation(gating_activation)
  
    def build(self, input_shape):
      num_reduced_filters = make_divisible(
          max(1, int(self._in_filters * self._se_ratio)),
          divisor=self._divisible_by)
  
      self._se_reduce = tf.keras.layers.Conv2D(
          filters=num_reduced_filters,
          kernel_size=1,
          strides=1,
          padding='same',
          use_bias=True,
          kernel_initializer=self._kernel_initializer,
          kernel_regularizer=self._kernel_regularizer,
          bias_regularizer=self._bias_regularizer)
  
      self._se_expand = tf.keras.layers.Conv2D(
          filters=self._out_filters,
          kernel_size=1,
          strides=1,
          padding='same',
          use_bias=True,
          kernel_initializer=self._kernel_initializer,
          kernel_regularizer=self._kernel_regularizer,
          bias_regularizer=self._bias_regularizer)
  
      super(SqueezeExcitation, self).build(input_shape)
  
    def get_config(self):
      config = {
          'in_filters': self._in_filters,
          'out_filters': self._out_filters,
          'se_ratio': self._se_ratio,
          'divisible_by': self._divisible_by,
          'use_3d_input': self._use_3d_input,
          'kernel_initializer': self._kernel_initializer,
          'kernel_regularizer': self._kernel_regularizer,
          'bias_regularizer': self._bias_regularizer,
          'activation': self._activation,
          'gating_activation': self._gating_activation,
      }
      base_config = super(SqueezeExcitation, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
    def call(self, inputs):
      x = tf.reduce_mean(inputs, self._spatial_axis, keepdims=True)
      x = self._activation_fn(self._se_reduce(x))
      x = self._gating_activation_fn(self._se_expand(x))
      return x * inputs
  
  
  def get_stochastic_depth_rate(init_rate, i, n):
    """Get drop connect rate for the ith block.
  
    Args:
      init_rate: A `float` of initial drop rate.
      i: An `int` of order of the current block.
      n: An `int` total number of blocks.
  
    Returns:
      Drop rate of the ith block.
    """
    if init_rate is not None:
      if init_rate < 0 or init_rate > 1:
        raise ValueError('Initial drop rate must be within 0 and 1.')
      rate = init_rate * float(i) / n
    else:
      rate = None
    return rate
  
  
  @tf.keras.utils.register_keras_serializable(package='Vision')
  class StochasticDepth(tf.keras.layers.Layer):
    """Creates a stochastic depth layer."""
  
    def __init__(self, stochastic_depth_drop_rate, **kwargs):
      """Initializes a stochastic depth layer.
  
      Args:
        stochastic_depth_drop_rate: A `float` of drop rate.
        **kwargs: Additional keyword arguments to be passed.
  
      Returns:
        A output `tf.Tensor` of which should have the same shape as input.
      """
      super(StochasticDepth, self).__init__(**kwargs)
      self._drop_rate = stochastic_depth_drop_rate
  
    def get_config(self):
      config = {'drop_rate': self._drop_rate}
      base_config = super(StochasticDepth, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
    def call(self, inputs, training=None):
      if training is None:
        training = tf.keras.backend.learning_phase()
      if not training or self._drop_rate is None or self._drop_rate == 0:
        return inputs
  
      keep_prob = 1.0 - self._drop_rate
      batch_size = tf.shape(inputs)[0]
      random_tensor = keep_prob
      random_tensor += tf.random.uniform(
          [batch_size] + [1] * (inputs.shape.rank - 1), dtype=inputs.dtype)
      binary_tensor = tf.floor(random_tensor)
      output = tf.math.divide(inputs, keep_prob) * binary_tensor
      return output
  
  
  @tf.keras.utils.register_keras_serializable(package='Vision')
  def pyramid_feature_fusion(inputs, target_level):
    """Fuses all feature maps in the feature pyramid at the target level.
  
    Args:
      inputs: A dictionary containing the feature pyramid. The size of the input
        tensor needs to be fixed.
      target_level: An `int` of the target feature level for feature fusion.
  
    Returns:
      A `float` `tf.Tensor` of shape [batch_size, feature_height, feature_width,
        feature_channel].
    """
    # Convert keys to int.
    pyramid_feats = {int(k): v for k, v in inputs.items()}
    min_level = min(pyramid_feats.keys())
    max_level = max(pyramid_feats.keys())
    resampled_feats = []
  
    for l in range(min_level, max_level + 1):
      if l == target_level:
        resampled_feats.append(pyramid_feats[l])
      else:
        feat = pyramid_feats[l]
        target_size = list(feat.shape[1:3])
        target_size[0] *= 2**(l - target_level)
        target_size[1] *= 2**(l - target_level)
        # Casts feat to float32 so the resize op can be run on TPU.
        feat = tf.cast(feat, tf.float32)
        feat = tf.image.resize(
            feat, size=target_size, method=tf.image.ResizeMethod.BILINEAR)
        # Casts it back to be compatible with the rest opetations.
        feat = tf.cast(feat, pyramid_feats[l].dtype)
        resampled_feats.append(feat)
  
    return tf.math.add_n(resampled_feats)
  
  
  @tf.keras.utils.register_keras_serializable(package='Vision')
  class Scale(tf.keras.layers.Layer):
    """Scales the input by a trainable scalar weight.
  
    This is useful for applying ReZero to layers, which improves convergence
    speed. This implements the paper:
    ReZero is All You Need: Fast Convergence at Large Depth.
    (https://arxiv.org/pdf/2003.04887.pdf).
    """
  
    def __init__(
        self,
        initializer: tf.keras.initializers.Initializer = 'ones',
        regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
        **kwargs):
      """Initializes a scale layer.
  
      Args:
        initializer: A `str` of initializer for the scalar weight.
        regularizer: A `tf.keras.regularizers.Regularizer` for the scalar weight.
        **kwargs: Additional keyword arguments to be passed to this layer.
  
      Returns:
        An `tf.Tensor` of which should have the same shape as input.
      """
      super(Scale, self).__init__(**kwargs)
  
      self._initializer = initializer
      self._regularizer = regularizer
  
      self._scale = self.add_weight(
          name='scale',
          shape=[],
          dtype=self.dtype,
          initializer=self._initializer,
          regularizer=self._regularizer,
          trainable=True)
  
    def get_config(self):
      """Returns a dictionary containing the config used for initialization."""
      config = {
          'initializer': self._initializer,
          'regularizer': self._regularizer,
      }
      base_config = super(Scale, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
    def call(self, inputs):
      """Calls the layer with the given inputs."""
      scale = tf.cast(self._scale, inputs.dtype)
      return scale * inputs
  
  
  @tf.keras.utils.register_keras_serializable(package='Vision')
  class TemporalSoftmaxPool(tf.keras.layers.Layer):
    """Creates a network layer corresponding to temporal softmax pooling.
  
    This is useful for multi-class logits (used in e.g., Charades). Modified from
    AssembleNet Charades evaluation from:
  
    Michael S. Ryoo, AJ Piergiovanni, Mingxing Tan, Anelia Angelova.
    AssembleNet: Searching for Multi-Stream Neural Connectivity in Video
    Architectures.
    (https://arxiv.org/pdf/1905.13209.pdf).
    """
  
    def call(self, inputs):
      """Calls the layer with the given inputs."""
      assert inputs.shape.rank in (3, 4, 5)
      frames = tf.shape(inputs)[1]
      pre_logits = inputs / tf.sqrt(tf.cast(frames, inputs.dtype))
      activations = tf.nn.softmax(pre_logits, axis=1)
      outputs = inputs * activations
      return outputs
  
  
  @tf.keras.utils.register_keras_serializable(package='Vision')
  class PositionalEncoding(tf.keras.layers.Layer):
    """Creates a network layer that adds a sinusoidal positional encoding.
  
    Positional encoding is incremented across frames, and is added to the input.
    The positional encoding is first weighted at 0 so that the network can choose
    to ignore it. This implements:
  
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones,
    Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin.
    Attention Is All You Need.
    (https://arxiv.org/pdf/1706.03762.pdf).
    """
  
    def __init__(self,
                 initializer: tf.keras.initializers.Initializer = 'zeros',
                 cache_encoding: bool = False,
                 state_prefix: Optional[str] = None,
                 **kwargs):
      """Initializes positional encoding.
  
      Args:
        initializer: A `str` of initializer for weighting the positional encoding.
        cache_encoding: A `bool`. If True, cache the positional encoding tensor
          after calling build. Otherwise, rebuild the tensor for every call.
          Setting this to False can be useful when we want to input a variable
          number of frames, so the positional encoding tensor can change shape.
        state_prefix: a prefix string to identify states.
        **kwargs: Additional keyword arguments to be passed to this layer.
  
      Returns:
        A `tf.Tensor` of which should have the same shape as input.
      """
      super(PositionalEncoding, self).__init__(**kwargs)
      self._initializer = initializer
      self._cache_encoding = cache_encoding
      self._pos_encoding = None
      self._rezero = Scale(initializer=initializer, name='rezero')
      state_prefix = state_prefix if state_prefix is not None else ''
      self._state_prefix = state_prefix
      self._frame_count_name = f'{state_prefix}_pos_enc_frame_count'
  
    def get_config(self):
      """Returns a dictionary containing the config used for initialization."""
      config = {
          'initializer': self._initializer,
          'cache_encoding': self._cache_encoding,
          'state_prefix': self._state_prefix,
      }
      base_config = super(PositionalEncoding, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
    def _positional_encoding(self,
                             num_positions: Union[int, tf.Tensor],
                             hidden_size: Union[int, tf.Tensor],
                             start_position: Union[int, tf.Tensor] = 0,
                             dtype: str = 'float32') -> tf.Tensor:
      """Creates a sequence of sinusoidal positional encoding vectors.
  
      Args:
        num_positions: the total number of positions (frames).
        hidden_size: the number of channels used for the hidden vectors.
        start_position: the start position.
        dtype: the dtype of the output tensor.
  
      Returns:
        The positional encoding tensor with shape [num_positions, hidden_size].
      """
      if isinstance(start_position, tf.Tensor) and start_position.shape.rank == 1:
        start_position = start_position[0]
  
      # Calling `tf.range` with `dtype=tf.bfloat16` results in an error,
      # so we cast afterward.
      positions = tf.range(start_position, start_position + num_positions)
      positions = tf.cast(positions, dtype)[:, tf.newaxis]
      idx = tf.range(hidden_size)[tf.newaxis, :]
  
      power = tf.cast(2 * (idx // 2), dtype)
      power /= tf.cast(hidden_size, dtype)
      angles = 1. / tf.math.pow(10_000., power)
      radians = positions * angles
  
      sin = tf.math.sin(radians[:, 0::2])
      cos = tf.math.cos(radians[:, 1::2])
      pos_encoding = tf.concat([sin, cos], axis=-1)
  
      return pos_encoding
  
    def _get_pos_encoding(self,
                          input_shape: tf.Tensor,
                          frame_count: int = 0) -> tf.Tensor:
      """Calculates the positional encoding from the input shape.
  
      Args:
        input_shape: the shape of the input.
        frame_count: a count of frames that indicates the index of the first
          frame.
  
      Returns:
        The positional encoding tensor with shape [num_positions, hidden_size].
  
      """
      frames = input_shape[1]
      channels = input_shape[-1]
      pos_encoding = self._positional_encoding(
          frames, channels, start_position=frame_count, dtype=self.dtype)
      pos_encoding = tf.reshape(pos_encoding, [1, frames, 1, 1, channels])
      return pos_encoding
  
    def build(self, input_shape):
      """Builds the layer with the given input shape.
  
      Args:
        input_shape: The input shape.
  
      Raises:
        ValueError: If using 'channels_first' data format.
      """
      if tf.keras.backend.image_data_format() == 'channels_first':
        raise ValueError('"channels_first" mode is unsupported.')
  
      if self._cache_encoding:
        self._pos_encoding = self._get_pos_encoding(input_shape)
  
      super(PositionalEncoding, self).build(input_shape)
  
    def call(
        self,
        inputs: tf.Tensor,
        states: Optional[States] = None,
        output_states: bool = True,
    ) -> Union[tf.Tensor, Tuple[tf.Tensor, States]]:
      """Calls the layer with the given inputs.
  
      Args:
        inputs: An input `tf.Tensor`.
        states: A `dict` of states such that, if any of the keys match for this
          layer, will overwrite the contents of the buffer(s). Expected keys
          include `state_prefix + '_pos_enc_frame_count'`.
        output_states: A `bool`. If True, returns the output tensor and output
          states. Returns just the output tensor otherwise.
  
      Returns:
        An output `tf.Tensor` (and optionally the states if `output_states=True`).
  
      Raises:
        ValueError: If using 'channels_first' data format.
      """
      states = dict(states) if states is not None else {}
  
      # Keep a count of frames encountered across input iterations in
      # num_frames to be able to accurately update the positional encoding.
      num_frames = tf.shape(inputs)[1]
      frame_count = tf.cast(states.get(self._frame_count_name, [0]), tf.int32)
      states[self._frame_count_name] = frame_count + num_frames
  
      if self._cache_encoding:
        pos_encoding = self._pos_encoding
      else:
        pos_encoding = self._get_pos_encoding(
            tf.shape(inputs), frame_count=frame_count)
      pos_encoding = tf.cast(pos_encoding, inputs.dtype)
      pos_encoding = self._rezero(pos_encoding)
      outputs = inputs + pos_encoding
  
      return (outputs, states) if output_states else outputs
  
  
  @tf.keras.utils.register_keras_serializable(package='Vision')
  class GlobalAveragePool3D(tf.keras.layers.Layer):
    """Creates a global average pooling layer with causal mode.
  
    Implements causal mode, which runs a cumulative sum (with `tf.cumsum`) across
    frames in the time dimension, allowing the use of a stream buffer. Sums any
    valid input state with the current input to allow state to accumulate over
    several iterations.
    """
  
    def __init__(self,
                 keepdims: bool = False,
                 causal: bool = False,
                 state_prefix: Optional[str] = None,
                 **kwargs):
      """Initializes a global average pool layer.
  
      Args:
        keepdims: A `bool`. If True, keep the averaged dimensions.
        causal: A `bool` of whether to run in causal mode with a cumulative sum
          across frames.
        state_prefix: a prefix string to identify states.
        **kwargs: Additional keyword arguments to be passed to this layer.
  
      Returns:
        An output `tf.Tensor`.
      """
      super(GlobalAveragePool3D, self).__init__(**kwargs)
  
      self._keepdims = keepdims
      self._causal = causal
      state_prefix = state_prefix if state_prefix is not None else ''
      self._state_prefix = state_prefix
  
      self._state_name = f'{state_prefix}_pool_buffer'
      self._frame_count_name = f'{state_prefix}_pool_frame_count'
  
    def get_config(self):
      """Returns a dictionary containing the config used for initialization."""
      config = {
          'keepdims': self._keepdims,
          'causal': self._causal,
          'state_prefix': self._state_prefix,
      }
      base_config = super(GlobalAveragePool3D, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
    def call(self,
             inputs: tf.Tensor,
             states: Optional[States] = None,
             output_states: bool = True
             ) -> Union[tf.Tensor, Tuple[tf.Tensor, States]]:
      """Calls the layer with the given inputs.
  
      Args:
        inputs: An input `tf.Tensor`.
        states: A `dict` of states such that, if any of the keys match for this
          layer, will overwrite the contents of the buffer(s).
          Expected keys include `state_prefix + '__pool_buffer'` and
          `state_prefix + '__pool_frame_count'`.
        output_states: A `bool`. If True, returns the output tensor and output
          states. Returns just the output tensor otherwise.
  
      Returns:
        An output `tf.Tensor` (and optionally the states if `output_states=True`).
        If `causal=True`, the output tensor will have shape
        `[batch_size, num_frames, 1, 1, channels]` if `keepdims=True`. We keep
        the frame dimension in this case to simulate a cumulative global average
        as if we are inputting one frame at a time. If `causal=False`, the output
        is equivalent to `tf.keras.layers.GlobalAveragePooling3D` with shape
        `[batch_size, 1, 1, 1, channels]` if `keepdims=True` (plus the optional
        buffer stored in `states`).
  
      Raises:
        ValueError: If using 'channels_first' data format.
      """
      states = dict(states) if states is not None else {}
  
      if tf.keras.backend.image_data_format() == 'channels_first':
        raise ValueError('"channels_first" mode is unsupported.')
  
      # Shape: [batch_size, 1, 1, 1, channels]
      buffer = states.get(self._state_name, None)
      if buffer is None:
        buffer = tf.zeros_like(inputs[:, :1, :1, :1], dtype=inputs.dtype)
        states[self._state_name] = buffer
  
      # Keep a count of frames encountered across input iterations in
      # num_frames to be able to accurately take a cumulative average across
      # all frames when running in streaming mode
      num_frames = tf.shape(inputs)[1]
      frame_count = states.get(self._frame_count_name, tf.constant([0]))
      frame_count = tf.cast(frame_count, tf.int32)
      states[self._frame_count_name] = frame_count + num_frames
  
      if self._causal:
        # Take a mean of spatial dimensions to make computation more efficient.
        x = tf.reduce_mean(inputs, axis=[2, 3], keepdims=True)
        x = tf.cumsum(x, axis=1)
        x = x + buffer
  
        # The last frame will be the value of the next state
        # Shape: [batch_size, 1, 1, 1, channels]
        states[self._state_name] = x[:, -1:]
  
        # In causal mode, the divisor increments by 1 for every frame to
        # calculate cumulative averages instead of one global average
        mean_divisors = tf.range(num_frames) + frame_count + 1
        mean_divisors = tf.reshape(mean_divisors, [1, num_frames, 1, 1, 1])
        mean_divisors = tf.cast(mean_divisors, x.dtype)
  
        # Shape: [batch_size, num_frames, 1, 1, channels]
        x = x / mean_divisors
      else:
        # In non-causal mode, we (optionally) sum across frames to take a
        # cumulative average across input iterations rather than individual
        # frames. If no buffer state is passed, this essentially becomes
        # regular global average pooling.
        # Shape: [batch_size, 1, 1, 1, channels]
        x = tf.reduce_sum(inputs, axis=(1, 2, 3), keepdims=True)
        x = x / tf.cast(tf.shape(inputs)[2] * tf.shape(inputs)[3], x.dtype)
        x = x + buffer
  
        # Shape: [batch_size, 1, 1, 1, channels]
        states[self._state_name] = x
  
        x = x / tf.cast(frame_count + num_frames, x.dtype)
  
      if not self._keepdims:
        x = tf.squeeze(x, axis=(1, 2, 3))
  
      return (x, states) if output_states else x
  
  
  @tf.keras.utils.register_keras_serializable(package='Vision')
  class SpatialAveragePool3D(tf.keras.layers.Layer):
    """Creates a global average pooling layer pooling across spatial dimentions."""
  
    def __init__(self, keepdims: bool = False, **kwargs):
      """Initializes a global average pool layer.
  
      Args:
        keepdims: A `bool`. If True, keep the averaged dimensions.
        **kwargs: Additional keyword arguments to be passed to this layer.
  
      Returns:
        An output `tf.Tensor`.
      """
      super(SpatialAveragePool3D, self).__init__(**kwargs)
      self._keepdims = keepdims
  
    def get_config(self):
      """Returns a dictionary containing the config used for initialization."""
      config = {
          'keepdims': self._keepdims,
      }
      base_config = super(SpatialAveragePool3D, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
    def build(self, input_shape):
      """Builds the layer with the given input shape."""
      if tf.keras.backend.image_data_format() == 'channels_first':
        raise ValueError('"channels_first" mode is unsupported.')
  
      super(SpatialAveragePool3D, self).build(input_shape)
  
    def call(self, inputs):
      """Calls the layer with the given inputs."""
      if inputs.shape.rank != 5:
        raise ValueError(
            'Input should have rank {}, got {}'.format(5, inputs.shape.rank))
  
      return tf.reduce_mean(inputs, axis=(2, 3), keepdims=self._keepdims)
  
  
  class CausalConvMixin:
    """Mixin class to implement CausalConv for `tf.keras.layers.Conv` layers."""
  
    @property
    def use_buffered_input(self) -> bool:
      return self._use_buffered_input
  
    @use_buffered_input.setter
    def use_buffered_input(self, variable: bool):
      self._use_buffered_input = variable
  
    def _compute_buffered_causal_padding(self,
                                         inputs: tf.Tensor,
                                         use_buffered_input: bool = False,
                                         time_axis: int = 1,
                                         ) -> List[List[int]]:
      """Calculates padding for 'causal' option for conv layers.
  
      Args:
        inputs: An optional input `tf.Tensor` to be padded.
        use_buffered_input: A `bool`. If True, use 'valid' padding along the time
          dimension. This should be set when applying the stream buffer.
        time_axis: An `int` of the axis of the time dimension.
  
      Returns:
        A list of paddings for `tf.pad`.
      """
      input_shape = tf.shape(inputs)[1:-1]
  
      if tf.keras.backend.image_data_format() == 'channels_first':
        raise ValueError('"channels_first" mode is unsupported.')
  
      kernel_size_effective = [
          (self.kernel_size[i] +
           (self.kernel_size[i] - 1) * (self.dilation_rate[i] - 1))
          for i in range(self.rank)
      ]
      pad_total = [kernel_size_effective[0] - 1]
      for i in range(1, self.rank):
        overlap = (input_shape[i] - 1) % self.strides[i] + 1
        pad_total.append(tf.maximum(kernel_size_effective[i] - overlap, 0))
      pad_beg = [pad_total[i] // 2 for i in range(self.rank)]
      pad_end = [pad_total[i] - pad_beg[i] for i in range(self.rank)]
      padding = [[pad_beg[i], pad_end[i]] for i in range(self.rank)]
      padding = [[0, 0]] + padding + [[0, 0]]
  
      if use_buffered_input:
        padding[time_axis] = [0, 0]
      else:
        padding[time_axis] = [padding[time_axis][0] + padding[time_axis][1], 0]
      return padding
  
    def _causal_validate_init(self):
      """Validates the Conv layer initial configuration."""
      # Overriding this method is meant to circumvent unnecessary errors when
      # using causal padding.
      if (self.filters is not None
          and self.filters % self.groups != 0):
        raise ValueError(
            'The number of filters must be evenly divisible by the number of '
            'groups. Received: groups={}, filters={}'.format(
                self.groups, self.filters))
  
      if not all(self.kernel_size):
        raise ValueError('The argument `kernel_size` cannot contain 0(s). '
                         'Received: %s' % (self.kernel_size,))
  
    def _buffered_spatial_output_shape(self, spatial_output_shape: List[int]):
      """Computes the spatial output shape from the input shape."""
      # When buffer padding, use 'valid' padding across time. The output shape
      # across time should be the input shape minus any padding, assuming
      # the stride across time is 1.
      if self._use_buffered_input and spatial_output_shape[0] is not None:
        padding = self._compute_buffered_causal_padding(
            tf.zeros([1] + spatial_output_shape + [1]), use_buffered_input=False)
        spatial_output_shape[0] -= sum(padding[1])
      return spatial_output_shape
  
  
  @tf.keras.utils.register_keras_serializable(package='Vision')
  class Conv2D(tf.keras.layers.Conv2D, CausalConvMixin):
    """Conv2D layer supporting CausalConv.
  
    Supports `padding='causal'` option (like in `tf.keras.layers.Conv1D`),
    which applies causal padding to the temporal dimension, and same padding in
    the spatial dimensions.
    """
  
    def __init__(self, *args, use_buffered_input=False, **kwargs):
      """Initializes conv2d.
  
      Args:
        *args: Arguments to be passed.
        use_buffered_input: A `bool`. If True, the input is expected to be padded
          beforehand. In effect, calling this layer will use 'valid' padding on
          the temporal dimension to simulate 'causal' padding.
        **kwargs: Additional keyword arguments to be passed.
  
      Returns:
        An output `tf.Tensor` of the Conv2D operation.
      """
      super(Conv2D, self).__init__(*args, **kwargs)
      self._use_buffered_input = use_buffered_input
  
    def get_config(self):
      """Returns a dictionary containing the config used for initialization."""
      config = {
          'use_buffered_input': self._use_buffered_input,
      }
      base_config = super(Conv2D, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
    def _compute_causal_padding(self, inputs):
      """Computes causal padding dimensions for the given inputs."""
      return self._compute_buffered_causal_padding(
          inputs, use_buffered_input=self._use_buffered_input)
  
    def _validate_init(self):
      """Validates the Conv layer initial configuration."""
      self._causal_validate_init()
  
    def _spatial_output_shape(self, spatial_input_shape: List[int]):
      """Computes the spatial output shape from the input shape."""
      shape = super(Conv2D, self)._spatial_output_shape(spatial_input_shape)
      return self._buffered_spatial_output_shape(shape)
  
  
  @tf.keras.utils.register_keras_serializable(package='Vision')
  class DepthwiseConv2D(tf.keras.layers.DepthwiseConv2D, CausalConvMixin):
    """DepthwiseConv2D layer supporting CausalConv.
  
    Supports `padding='causal'` option (like in `tf.keras.layers.Conv1D`),
    which applies causal padding to the temporal dimension, and same padding in
    the spatial dimensions.
    """
  
    def __init__(self, *args, use_buffered_input=False, **kwargs):
      """Initializes depthwise conv2d.
  
      Args:
        *args: Arguments to be passed.
        use_buffered_input: A `bool`. If True, the input is expected to be padded
          beforehand. In effect, calling this layer will use 'valid' padding on
          the temporal dimension to simulate 'causal' padding.
        **kwargs: Additional keyword arguments to be passed.
  
      Returns:
        An output `tf.Tensor` of the DepthwiseConv2D operation.
      """
      super(DepthwiseConv2D, self).__init__(*args, **kwargs)
      self._use_buffered_input = use_buffered_input
  
      # Causal padding is unsupported by default for DepthwiseConv2D,
      # so we resort to valid padding internally. However, we handle
      # causal padding as a special case with `self._is_causal`, which is
      # defined by the super class.
      if self.padding == 'causal':
        self.padding = 'valid'
  
    def get_config(self):
      """Returns a dictionary containing the config used for initialization."""
      config = {
          'use_buffered_input': self._use_buffered_input,
      }
      base_config = super(DepthwiseConv2D, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
    def call(self, inputs):
      """Calls the layer with the given inputs."""
      if self._is_causal:
        inputs = tf.pad(inputs, self._compute_causal_padding(inputs))
      return super(DepthwiseConv2D, self).call(inputs)
  
    def _compute_causal_padding(self, inputs):
      """Computes causal padding dimensions for the given inputs."""
      return self._compute_buffered_causal_padding(
          inputs, use_buffered_input=self._use_buffered_input)
  
    def _validate_init(self):
      """Validates the Conv layer initial configuration."""
      self._causal_validate_init()
  
    def _spatial_output_shape(self, spatial_input_shape: List[int]):
      """Computes the spatial output shape from the input shape."""
      shape = super(DepthwiseConv2D, self)._spatial_output_shape(
          spatial_input_shape)
      return self._buffered_spatial_output_shape(shape)
  
  
  @tf.keras.utils.register_keras_serializable(package='Vision')
  class Conv3D(tf.keras.layers.Conv3D, CausalConvMixin):
    """Conv3D layer supporting CausalConv.
  
    Supports `padding='causal'` option (like in `tf.keras.layers.Conv1D`),
    which applies causal padding to the temporal dimension, and same padding in
    the spatial dimensions.
    """
  
    def __init__(self, *args, use_buffered_input=False, **kwargs):
      """Initializes conv3d.
  
      Args:
        *args: Arguments to be passed.
        use_buffered_input: A `bool`. If True, the input is expected to be padded
          beforehand. In effect, calling this layer will use 'valid' padding on
          the temporal dimension to simulate 'causal' padding.
        **kwargs: Additional keyword arguments to be passed.
  
      Returns:
        An output `tf.Tensor` of the Conv3D operation.
      """
      super(Conv3D, self).__init__(*args, **kwargs)
      self._use_buffered_input = use_buffered_input
  
    def get_config(self):
      """Returns a dictionary containing the config used for initialization."""
      config = {
          'use_buffered_input': self._use_buffered_input,
      }
      base_config = super(Conv3D, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
    def call(self, inputs):
      """Call the layer with the given inputs."""
      # Note: tf.nn.conv3d with depthwise kernels on CPU is currently only
      # supported when compiling with TF graph (XLA) using tf.function, so it
      # is compiled by default here (b/186463870).
      conv_fn = tf.function(super(Conv3D, self).call, jit_compile=True)
      return conv_fn(inputs)
  
    def _compute_causal_padding(self, inputs):
      """Computes causal padding dimensions for the given inputs."""
      return self._compute_buffered_causal_padding(
          inputs, use_buffered_input=self._use_buffered_input)
  
    def _validate_init(self):
      """Validates the Conv layer initial configuration."""
      self._causal_validate_init()
  
    def _spatial_output_shape(self, spatial_input_shape: List[int]):
      """Computes the spatial output shape from the input shape."""
      shape = super(Conv3D, self)._spatial_output_shape(spatial_input_shape)
      return self._buffered_spatial_output_shape(shape)
  