How to do atrous convolution in tensorflow 2 (tf.keras)

I am trying to convert some code from tensorflow 1.x to tensorflow 2.x. It's been going well so far, but I'm stuck on atrous convolution. Unlike other layers, there doesn't seem to be a one-to-one conversion.

So far, I've been unifying everything to tf.keras. There is a pure keras implementation here and a tf.nn.atrous_conv2d implementation here but I'm also not certain if I can use them in a tf.keras.Model functional api.

Here is the code:

with tf.variable_scope('aconv1d_' + name):
        shape = [None, 30, 128]
        kernel = tf.get_variable('kernel', (1, size, shape[-1], n_filters), dtype=tf.float32,
                                 initializer=tf.contrib.layers.xavier_initializer())
        if bias:
            b = tf.get_variable('b', [shape[-1]], dtype=tf.float32, initializer=tf.constant_initializer(0))
        out = tf.nn.atrous_conv2d(tf.expand_dims(input_tensor, dim=1), kernel, rate=rate, padding='SAME') + (
            b if bias else 0)
        out = tf.squeeze(out, [1])

        return out

I would like to just convert this, stick it in the keras functional api, do model.fit, and run.

Thank you for helping out a noob like me.