<gh_stars>100-1000
""" Mixins for nn.Modules for better textual visualization. """
from textwrap import indent



class LayerReprMixin:
    """ Adds useful properties and methods for nn.Modules, mainly related to visualization and introspection. """
    VERBOSITY_THRESHOLD = 10

    @property
    def num_frozen_parameters(self):
        return sum(p.numel() for p in self.parameters() if not p.requires_grad)

    @property
    def num_trainable_parameters(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    @property
    def num_parameters(self):
        return sum(p.numel() for p in self.parameters())

    def __repr__(self):
        if hasattr(self, 'verbosity') and self.verbosity < self.VERBOSITY_THRESHOLD:
            return ''

        msg = super().__repr__()
        if getattr(self, 'collapsible', False):
            msg = msg.replace('(\n  (layer): ', ':').replace('\n    ', '\n  ')
            msg = msg.replace('\n  )\n)', '\n)').replace(')\n)', ')')
        return msg

    def prepare_repr(self, verbosity=1, collapsible=True, show_num_parameters=False, extra=False):
        """ Set flags on children, call `repr`, delete flags.
        Returns string.
        """
        def set_flags(module):
            setattr(module, 'verbosity', verbosity)
            setattr(module, 'collapsible', collapsible)
            setattr(module, 'show_num_parameters', show_num_parameters)
            setattr(module, 'extra', extra)

        def del_flags(module):
            try:
                delattr(module, 'verbosity')
                delattr(module, 'collapsible')
                delattr(module, 'show_num_parameters')
                delattr(module, 'extra')
            except AttributeError:
                pass

        self.apply(set_flags)
        msg = repr(self)
        self.apply(del_flags)
        return msg

    def repr(self, verbosity=1, collapsible=True, show_num_parameters=False, extra=False):
        """ Set flags on children, call `repr`, delete flags.
        Prints output to stdout.
        """
        print(self.prepare_repr(verbosity=verbosity, collapsible=collapsible,
                                show_num_parameters=show_num_parameters, extra=extra))



class ModuleDictReprMixin(LayerReprMixin):
    """ Mixin to allow `repr` for multiple levels for nn.ModuleDicts.
    Also adds `__getitem__` for convenience.

    Relies on modules having `shapes` dictionary, that for each children stores the information about their
    input and output shapes.

    Depending on `verbosity`, creates string representation for different levels:
        - verbosity 1, modules and their shapes. For example, shapes of `initial_block`, `body` and `head`.
        - verbosity 2, blocks inside modules. For example, shapes of blocks inside Encoder.
        - verbosity 3, blocks inside repeated chains of `:class:~.blocks.Block`. Mainly used for debug purposes.
        - verbosity 4, letters inside each `:class:~.layers.MultiLayer`. For example, each letter inside given block.
        - verbosity 5, PyTorch implementations of each letter inside `:class:~.layers.MultiLayer`.
        - verbosity 6+, default repr of nn.Module.
    For most cases, levels 2 and 4 should be used.

    Additional parameters can be used to show number of parameters inside each level and collapse multilines.
    """
    def __getitem__(self, key):
        if isinstance(key, int):
            key = list(self.keys())[key]
        return super().__getitem__(key)

    def prepare_shape(self, shape, indent=0):
        """ Beautify shape or list of shapes.
        Changes the first dimension (batch) to `?`.
        Makes multiple lines for lists of shapes with provided indentation.
        """
        #pylint: disable=redefined-outer-name
        if isinstance(shape, tuple):
            msg = ', '.join([f'{item:>3}' for item in shape[1:]])
            return f' (?, {msg}) '

        if isinstance(shape, list):
            msg = '[' + self.prepare_shape(shape[0])[1:-1] + ','
            for shape_ in shape[1:]:
                msg += '\n ' + ' '*indent + self.prepare_shape(shape_)[1:-1] + ','
            msg = msg[:-1] + ']'
            return msg
        raise TypeError(f'Should be used on tuple or list of tuples, got {type(shape)} instead.')


    def __repr__(self):
        if hasattr(self, 'verbosity'):
            indent_prefix = '    '

            # Parse verbosity. If equal to max level, set flag
            verbosity = self.verbosity
            if verbosity >= 5:
                verbosity = 4
                detailed_last_level = True
            else:
                detailed_last_level = False

            if len(self.keys()):
                key = list(self.keys())[0]
                input_shapes, output_shapes = None, None

                if (len(self.items()) == 1 and getattr(self, 'collapsible', False)
                    and getattr(self[key], 'VERBOSITY_THRESHOLD', -1) == self.VERBOSITY_THRESHOLD):
                    # Subclasses names can be folded, i.e. `Block:ResBlock(` instead of `Block(\n    ResBlock(`
                    msg = self._get_name() + ':' + repr(self[key])
                    msg = msg.replace(')\n)', ')')

                else:
                    msg = self._get_name() + '(\n'
                    extra_repr = self.extra_repr()
                    if extra_repr:
                        msg += indent(extra_repr, prefix=indent_prefix) + '\n'

                    max_key_length = max(len(key) for key in self.keys())
                    for key, value in self.items():

                        # Short description: module name and description of shapes
                        empty_space = ' ' * (1 + max_key_length - len(key))
                        module_short_description = f'({key}:{empty_space}'

                        if key in self.shapes:
                            input_shapes, output_shapes = self.shapes.get(key)

                            current_line_len = len(module_short_description)
                            input_shapes = self.prepare_shape(input_shapes, indent=current_line_len)
                            module_short_description += input_shapes + ' ⟶ '

                            current_line_len = len(module_short_description.splitlines()[-1]) + 1
                            output_shapes = self.prepare_shape(output_shapes, indent=current_line_len).strip(' ')
                            module_short_description += output_shapes

                        if getattr(self, 'show_num_parameters', False):
                            num_parameters = sum(p.numel() for p in value.parameters() if p.requires_grad)
                            module_short_description += f', #params={num_parameters:,}'
                        module_short_description += ')'

                        # Long description: ~unmodified repr of a module
                        module_long_description = repr(value).strip(' ')

                        # Select appropriate message
                        module_description = ''
                        if verbosity > self.VERBOSITY_THRESHOLD:
                            module_description = f'({key}): ' + module_long_description

                        if verbosity == self.VERBOSITY_THRESHOLD or module_description == f'({key}): ':
                            module_description = module_short_description

                        if self.VERBOSITY_THRESHOLD == 4 and detailed_last_level:
                            module_description = (module_short_description + ':\n' +
                                                  indent(module_long_description, prefix=indent_prefix))

                        msg += indent(module_description, prefix=indent_prefix) + '\n'
                        msg = msg.replace('\n\n', '\n')
                    msg += ')'

                if len(self.items()) == 1 and getattr(self, 'collapsible', False) and 'r0' in msg:
                    msg = msg.replace('(\n    (r0:', '(r0:', 1)
                    msg = msg.replace(')\n)', ')')
            else:
                msg = self._get_name() + '()'
            return msg
        return super().__repr__()


REPR_DOC = '\n'.join(ModuleDictReprMixin.__doc__.split('\n')[3:])
LayerReprMixin.repr.__doc__ += '\n' + REPR_DOC
LayerReprMixin.prepare_repr.__doc__ += '\n' + REPR_DOC

""" Mixins for nn.Modules for better textual visualization. """
from textwrap import indent



class LayerReprMixin:
    """ Adds useful properties and methods for nn.Modules, mainly related to visualization and introspection. """
    VERBOSITY_THRESHOLD = 10

    @property
    def num_frozen_parameters(self):
        return sum(p.numel() for p in self.parameters() if not p.requires_grad)

    @property
    def num_trainable_parameters(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    @property
    def num_parameters(self):
        return sum(p.numel() for p in self.parameters())

    def __repr__(self):
        if hasattr(self, 'verbosity') and self.verbosity < self.VERBOSITY_THRESHOLD:
            return ''

        msg = super().__repr__()
        if getattr(self, 'collapsible', False):
            msg = msg.replace('(\n  (layer): ', ':').replace('\n    ', '\n  ')
            msg = msg.replace('\n  )\n)', '\n)').replace(')\n)', ')')
        return msg

    def prepare_repr(self, verbosity=1, collapsible=True, show_num_parameters=False, extra=False):
        """ Set flags on children, call `repr`, delete flags.
        Returns string.
        """
        def set_flags(module):
            setattr(module, 'verbosity', verbosity)
            setattr(module, 'collapsible', collapsible)
            setattr(module, 'show_num_parameters', show_num_parameters)
            setattr(module, 'extra', extra)

        def del_flags(module):
            try:
                delattr(module, 'verbosity')
                delattr(module, 'collapsible')
                delattr(module, 'show_num_parameters')
                delattr(module, 'extra')
            except AttributeError:
                pass

        self.apply(set_flags)
        msg = repr(self)
        self.apply(del_flags)
        return msg

    def repr(self, verbosity=1, collapsible=True, show_num_parameters=False, extra=False):
        """ Set flags on children, call `repr`, delete flags.
        Prints output to stdout.
        """
        print(self.prepare_repr(verbosity=verbosity, collapsible=collapsible,
                                show_num_parameters=show_num_parameters, extra=extra))



class ModuleDictReprMixin(LayerReprMixin):
    """ Mixin to allow `repr` for multiple levels for nn.ModuleDicts.
    Also adds `__getitem__` for convenience.

    Relies on modules having `shapes` dictionary, that for each children stores the information about their
    input and output shapes.

    Depending on `verbosity`, creates string representation for different levels:
        - verbosity 1, modules and their shapes. For example, shapes of `initial_block`, `body` and `head`.
        - verbosity 2, blocks inside modules. For example, shapes of blocks inside Encoder.
        - verbosity 3, blocks inside repeated chains of `:class:~.blocks.Block`. Mainly used for debug purposes.
        - verbosity 4, letters inside each `:class:~.layers.MultiLayer`. For example, each letter inside given block.
        - verbosity 5, PyTorch implementations of each letter inside `:class:~.layers.MultiLayer`.
        - verbosity 6+, default repr of nn.Module.
    For most cases, levels 2 and 4 should be used.

    Additional parameters can be used to show number of parameters inside each level and collapse multilines.
    """
    def __getitem__(self, key):
        if isinstance(key, int):
            key = list(self.keys())[key]
        return super().__getitem__(key)

    def prepare_shape(self, shape, indent=0):
        """ Beautify shape or list of shapes.
        Changes the first dimension (batch) to `?`.
        Makes multiple lines for lists of shapes with provided indentation.
        """
        #pylint: disable=redefined-outer-name
        if isinstance(shape, tuple):
            msg = ', '.join([f'{item:>3}' for item in shape[1:]])
            return f' (?, {msg}) '

        if isinstance(shape, list):
            msg = '[' + self.prepare_shape(shape[0])[1:-1] + ','
            for shape_ in shape[1:]:
                msg += '\n ' + ' '*indent + self.prepare_shape(shape_)[1:-1] + ','
            msg = msg[:-1] + ']'
            return msg
        raise TypeError(f'Should be used on tuple or list of tuples, got {type(shape)} instead.')


    def __repr__(self):
        if hasattr(self, 'verbosity'):
            indent_prefix = '    '

            # Parse verbosity. If equal to max level, set flag
            verbosity = self.verbosity
            if verbosity >= 5:
                verbosity = 4
                detailed_last_level = True
            else:
                detailed_last_level = False

            if len(self.keys()):
                key = list(self.keys())[0]
                input_shapes, output_shapes = None, None

                if (len(self.items()) == 1 and getattr(self, 'collapsible', False)
                    and getattr(self[key], 'VERBOSITY_THRESHOLD', -1) == self.VERBOSITY_THRESHOLD):
                    # Subclasses names can be folded, i.e. `Block:ResBlock(` instead of `Block(\n    ResBlock(`
                    msg = self._get_name() + ':' + repr(self[key])
                    msg = msg.replace(')\n)', ')')

                else:
                    msg = self._get_name() + '(\n'
                    extra_repr = self.extra_repr()
                    if extra_repr:
                        msg += indent(extra_repr, prefix=indent_prefix) + '\n'

                    max_key_length = max(len(key) for key in self.keys())
                    for key, value in self.items():

                        # Short description: module name and description of shapes
                        empty_space = ' ' * (1 + max_key_length - len(key))
                        module_short_description = f'({key}:{empty_space}'

                        if key in self.shapes:
                            input_shapes, output_shapes = self.shapes.get(key)

                            current_line_len = len(module_short_description)
                            input_shapes = self.prepare_shape(input_shapes, indent=current_line_len)
                            module_short_description += input_shapes + ' ⟶ '

                            current_line_len = len(module_short_description.splitlines()[-1]) + 1
                            output_shapes = self.prepare_shape(output_shapes, indent=current_line_len).strip(' ')
                            module_short_description += output_shapes

                        if getattr(self, 'show_num_parameters', False):
                            num_parameters = sum(p.numel() for p in value.parameters() if p.requires_grad)
                            module_short_description += f', #params={num_parameters:,}'
                        module_short_description += ')'

                        # Long description: ~unmodified repr of a module
                        module_long_description = repr(value).strip(' ')

                        # Select appropriate message
                        module_description = ''
                        if verbosity > self.VERBOSITY_THRESHOLD:
                            module_description = f'({key}): ' + module_long_description

                        if verbosity == self.VERBOSITY_THRESHOLD or module_description == f'({key}): ':
                            module_description = module_short_description

                        if self.VERBOSITY_THRESHOLD == 4 and detailed_last_level:
                            module_description = (module_short_description + ':\n' +
                                                  indent(module_long_description, prefix=indent_prefix))

                        msg += indent(module_description, prefix=indent_prefix) + '\n'
                        msg = msg.replace('\n\n', '\n')
                    msg += ')'

                if len(self.items()) == 1 and getattr(self, 'collapsible', False) and 'r0' in msg:
                    msg = msg.replace('(\n    (r0:', '(r0:', 1)
                    msg = msg.replace(')\n)', ')')
            else:
                msg = self._get_name() + '()'
            return msg
        return super().__repr__()


REPR_DOC = '\n'.join(ModuleDictReprMixin.__doc__.split('\n')[3:])
LayerReprMixin.repr.__doc__ += '\n' + REPR_DOC
LayerReprMixin.prepare_repr.__doc__ += '\n' + REPR_DOC
Mixins for nn.Modules for better textual visualization. Adds useful properties and methods for nn.Modules, mainly related to visualization and introspection. Set flags on children, call `repr`, delete flags.         Returns string. Set flags on children, call `repr`, delete flags.         Prints output to stdout. Mixin to allow `repr` for multiple levels for nn.ModuleDicts.     Also adds `__getitem__` for convenience.      Relies on modules having `shapes` dictionary, that for each children stores the information about their     input and output shapes.      Depending on `verbosity`, creates string representation for different levels:         - verbosity 1, modules and their shapes. For example, shapes of `initial_block`, `body` and `head`.         - verbosity 2, blocks inside modules. For example, shapes of blocks inside Encoder.         - verbosity 3, blocks inside repeated chains of `:class:~.blocks.Block`. Mainly used for debug purposes.         - verbosity 4, letters inside each `:class:~.layers.MultiLayer`. For example, each letter inside given block.         - verbosity 5, PyTorch implementations of each letter inside `:class:~.layers.MultiLayer`.         - verbosity 6+, default repr of nn.Module.     For most cases, levels 2 and 4 should be used.      Additional parameters can be used to show number of parameters inside each level and collapse multilines. Beautify shape or list of shapes.         Changes the first dimension (batch) to `?`.         Makes multiple lines for lists of shapes with provided indentation. #pylint: disable=redefined-outer-name # Parse verbosity. If equal to max level, set flag # Subclasses names can be folded, i.e. `Block:ResBlock(` instead of `Block(\n    ResBlock(` # Short description: module name and description of shapes #params={num_parameters:,}' # Long description: ~unmodified repr of a module # Select appropriate message