tf.keras.ops.space_to_depth

Rearranges blocks of spatial data into depth.

tf.keras.ops.space_to_depth(
    x, block_size, data_format='channels_last'
)

This operation is useful for resizing the activations between convolutions (but keeping all data). It is also useful for training purely convolutional models.

This operation rearranges elements in a tensor of shape (N, H * block_size, W * block_size, C) to a tensor of shape (N, H, W, C * block_size ** 2) (for data_format="channels_last") or (N, C, H * block_size, W * block_size) to (N, C * block_size ** 2, H, W) (for data_format="channels_first").

This is the reverse transformation of depth_to_space.

x Input tensor. Must be 4D.
block_size An integer specifying the size of the spatial block. The height and width of the input must be divisible by block_size.
data_format A string specifying the data format of the input tensor. "channels_last" corresponds to inputs with shape (batch, height, width, channels) while "channels_first" corresponds to inputs with shape (batch, channels, height, width). Defaults to "channels_last".

A tensor with the same dtype as x, with shape (N, H // block_size, W // block_size, C * block_size ** 2) for data_format="channels_last" or (N, C * block_size ** 2, H // block_size, W // block_size) for data_format="channels_first".

Example:

x = keras.ops.reshape(keras.ops.arange(1 * 4 * 4 * 3), (1, 4, 4, 3))
keras.ops.space_to_depth(x, block_size=2).shape
(1, 2, 2, 12)
 
# channels_first example
x = keras.ops.reshape(keras.ops.arange(1 * 3 * 4 * 4), (1, 3, 4, 4))
keras.ops.space_to_depth(x, block_size=2,
                         data_format="channels_first").shape
(1, 12, 2, 2)