import tensorflow as tf # Code from https://github.com/leanderme/ConvNeXt-Tensorflow/blob/main/ConvNeXt.ipynb class LayerStack2Image(tf.keras.layers.Layer): def __init__(self, target_width, target_height, name=None, **kwargs): super(LayerStack2Image, self).__init__(name=name) self.param_target_width = target_width self.param_target_height = target_height def get_config(self): config = super(LayerStack2Image, self).get_config() config.update({ "target_width": self.param_target_width, "target_height": self.param_target_height, }) return config def call(self, input_thing, **kwargs): result = tf.stack([ input_thing for i in range(self.param_target_width) ], axis=-1) result = tf.stack([ result for i in range(self.param_target_height) ], axis=-1) result = tf.stack([ result ], axis=-1) # channel dimension return result