Reflective padding as pure keras verision

Question

I am using keras with the plaidML backend and need to implement reflective padding.
With a tensorflow backend that is simply tf.pad with mode set to REFLECT.

How can implement that functionality with K. functions or plaidml tile functions ?
Or is there an implementation somewhere i could use ?

With K. function simply slicing, inverting and concatenating the value back together would probably be possible, but all my tries in that direction where a mess and didn't really worked out.

Answer 1

Here a version for reflective padding as pure K function, which should (but not tested) work with every backend:

def reflection_padding(inp, paddings):
    paddings = [(x, x) if isinstance(x, int) else x for x in paddings]
    ishape = inp.shape.dims
    ndims = inp.shape.ndims
    if len(ishape) != len(paddings):
        raise ValueError("Padding dims != input dims")
    last = inp
    _all_slice = slice(None, None, None)

    def _get_slices(ndims, axis, slice_):
        ret = [_all_slice for _ in range(ndims)]
        ret[axis] = slice_
        return tuple(ret)

    for axis, pads in ((i, x) for i, x in enumerate(paddings) if x[0]+x[1] != 0):
        pad_data = []
        if pads[0]:
            pre = last[_get_slices(ndims, axis, slice(pads[0], 0, -1))]
            pad_data.append(pre)
        pad_data.append(last)
        if pads[1]:
            post = last[_get_slices(ndims, axis, slice(-2, -pads[1]-2, -1))]
            pad_data.append(post)
        last = K.concatenate(pad_data, axis)
        ishape = last.shape.dims
    return last

# USAGE: reflection_padding(image_batch, [0, [2,2], [2,2], 0])

I am going to accept my own answer. If somebody has a better answer i'll gladly switch that over to theirs