How can I manipulate the Keras Layer ? : Stereo vision cost volume

I want to make a function that concatenates the Keras layers as following :

Suppose that we have two layers with same size of ( h x w x f )

What I want to do is, to make a feature volume ( w x h x w x f ) by following step :

I could not find a proper way to shift/cut layers by arbitrary number/axis.

If anyone can teach the way, thank you very much.

You can do the following. We are basically taking the slice as you have shown and then padding 0s to the end so that it has the same size as the blue input for stacking. Note that here, I'm doing the changes on height dimension (i.e. not width) so slight modification is required (mostly shuffling) to do it for width.

import tensorflow as tf
import tensorflow.keras.layers as layers

inp = layers.Input(shape=(64,64,3))

out1 = layers.Conv2D(16, (3,3), activation='relu', padding='same')(inp)
out2 = layers.Conv2D(16, (3,3), activation='relu', padding='same')(inp)

out3 = layers.Lambda(lambda x: tf.stack(
    [x[0]]+[
        tf.pad(x[1][:,:-i,:,:],[[0,0], [0,i],[0,0],[0,0]]) for i in range(1,x[1].shape[1])
    ], 
    axis=1))([out1, out2])

print(out3.shape)

which gives size of the final output as,

(None, 64, 64, 64, 16)