keras blurpool lambda

def get_blur_kernel(n_ch):
    filter_bank = np.zeros((5, 5, n_ch, 1), dtype="float32")
    v = np.array([1, 4, 6, 4, 1], dtype="float32")
    k = np.outer(v, v)
    k /= k.sum()
    for i in range(n_ch):
        filter_bank[:, :, i, 0] = k
    return tf.constant(filter_bank)

def blurpool(x):
    x_padded = tf.pad(x, [[0, 0], [2, 2], [2, 2], [0, 0]], "REFLECT")
    y = tf.nn.depthwise_conv2d_native(x_padded, get_blur_kernel(x.shape[-1]),
                                      strides=[1, 2, 2, 1], padding="VALID")
    return y


def test_model():
    x = tf.keras.layers.Input((32, 32, 3))
    y = tf.keras.layers.Conv2D(64, 3, padding="same")(x)
    y = tf.keras.layers.Lambda(blurpool)(y)
    return tf.keras.models.Model(x, y)