import tensorflow as tf

def _get_avgpool_scale(kw, kh):
  if kh > 255 or kw > 255:
    return 1.0
  elif kh == 3 and kw == 3:
    return 9.0 * 7.0 / 64.0
  elif kh == 5 and kw == 5:
    return 25.0 * 10.0 / 256.0
  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):
