from tensorflow.python.ops import array_ops

    super(VitisAveragePooling2D, self).build(input_shape)
    # Compute rescale factor in build() since the pool_size is determined.
    self.rescale_factor = _get_avgpool_scale(self.pool_size[0],
                                             self.pool_size[1])

  def call(self, inputs):
    outputs = super(VitisAveragePooling2D, self).call(inputs)

    # Simulate DPU hahavior of AvgPooling
    input_shape = array_ops.shape(inputs)

    if self.rescale_factor != 1.0:
      outputs *= self.rescale_factor
    return outputs


def _types_dict():
  return {
