[Python] keras model equivalent of tf.depth_to_space

I want to accomplish the equivalent of tf.depth_to_space in a Keras model. Specifically, the data in the Keras model is shaped H x W x 4 (i.e., depth of 4) and I want to permute the data so that the output is sized H x W x 1, with the mapping done as viewing the 4 input channels as 2x2 blocks; i.e.,

input location is y, x, k

output location is 2*y+(k//2), 2*x+(k%2), 1

I know that I can get the correct shape with:

outputs = keras.layers.Reshape((H*2,W*2,1), input_shape=(H,W,4))(inputs)


But I think that the mapping will be

input location is y, x, k

Linear_addess is y*W*4+x*4+k

output location is Linear_addess//(H*2), Linear_addess % (H*2), 1

which is not what I want

I tried directly using the

outputs = tf.depth_to_space(inputs, 2)


but that lead to an error: TypeError: Output tensors to a Model must be Keras tensors. Found Tensor("DepthToSpace:0", shape=(?, 1024, 1024, 1), dtype=float32)

the problem can be seen with this simple function

def simple_net(H=512, W=512):

inputs = keras.layers.Input((H, W, 4))

# gets the correct shape but not the correct order
outputs = keras.layers.Reshape((H*2,W*2,1), input_shape=(H,W,4))(inputs)

# Run time error message
#outputs = tf.depth_to_space(output_planes, 2)
model = keras.models.Model(inputs, outputs)
return model

Continue reading...