from tensorflow.python.keras.utils.generic_utils import register_keras_serializable

    ns_pow = tf.pow(2., ns)
    ks = tf.round(ns_pow / rec)
    diffs = tf.math.abs(ks / ns_pow - 1 / rec)
    n = tf.argmin(diffs)
    k = ks[n]
    scale = k / tf.pow(2., tf.cast(n, tf.float32))
    scale *= rec
    return scale


@register_keras_serializable(
    package='Vitis', name='VitisGlobalAveragePooling2D')
class VitisGlobalAveragePooling2D(tf.keras.layers.GlobalAveragePooling2D):
  """Vitis version of GlobalAveragePooling2D layer.

  This is an Vitis version of average pooling to simulate DPU behaviour which to
  integer approximations for averaging of specific sizes.
  """

  def __init__(self, **kwargs):
