import tensorflow as tf

  elif kh == 6 and kw == 6:
    return 36.0 * 7.0 / 256.0
  elif kh == 7 and kw == 7:
    return 49.0 * 21.0 / 1024.0
  elif kh == 14 and kw == 14:
    return 196.0 * 21.0 / 4096.0
  else:
    rec = tf.cast(kw * kh, tf.float32)
    n_max = 7 + tf.math.ceil(tf.math.log(rec) / tf.math.log(2.))
    ns = tf.range(0., n_max)
    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
