def _kernel_shape(self, input_shape):
    """Helper to calculate the kernel shape."""
    kernel_size_iter = iter(self._kernel_size)
    return [
        self._filters if c == 'O' else input_shape[self._lhs_spec.index('C')]
        if c == 'I' else next(kernel_size_iter) for c in self._rhs_spec
    ]
Helper to calculate the kernel shape.
trax/layers/deconvolution.py