Debug layer for trax, prints out the shape of the data it works with

debug_trax_layer.py

from trax.layers import base

class Debug(base.Layer):
    def __init__(self, msg=""):
        super().__init__(name=f'Debug')
        self.msg = msg
        self.debug = False

    def forward(self, x):
        if self.debug:
            print(f"{self.msg} {x.shape}")
        return x

    def init_weights_and_state(self, input_signature):
        pass

After training, turn on the debug mode with

model.sublayers[0].debug = True

(or whatever is appropriate for your model).