Neat CoordConv channels injection implementation as a tensorflow.keras layer.

coordconv.py

from tensorflow.keras import backend as K
from tensorflow.keras.layers import Layer


class CoordinateChannel2D(Layer):

    def call(self, inputs):
        
        x = K.cast(K.arange(0, K.shape(inputs)[1]), K.floatx())
        x /= K.cast(K.shape(inputs)[1], K.floatx())
        x = K.tile([[x]], [K.shape(inputs)[0], K.shape(inputs)[2], 1])
        x = K.expand_dims(x, -1)

        y = K.cast(K.arange(0, K.shape(inputs)[2]), K.floatx())
        y /= K.cast(K.shape(inputs)[2], K.floatx())
        y = K.tile([[y]], [K.shape(inputs)[0], K.shape(inputs)[1], 1])
        y = K.permute_dimensions(y, [0, 2, 1])
        y = K.expand_dims(y, -1)

        ret = K.concatenate([inputs, x, y], axis=-1)

        return ret