Repository: tensorflow/models
File: official/vision/beta/modeling/layers/nn_layers.py (Python)

Code before:
  """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: Optional[tf.Tensor] = None,
                                       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`.
    """
    shape_in = inputs.shape[1:-1]
    del inputs

    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 = [max(kernel_size_effective[i] - (self.strides[i]), 0)
                 if (shape_in[i]%self.strides[i]) == 0 else
                 max(kernel_size_effective[i] - 
                     (shape_in[i]%self.strides[i]), 0)
                 for i in range(self.rank)]
    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

Review comment (suggestion):
Let's rename to `input_shape`. Also, `inputs.shape` does not work in graph mode, so we should use `tf.shape(inputs)`:

```suggestion
    input_shape = tf.shape(inputs)[1:-1]
```

Code after:
  """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):