def _create_weight_decay_optim_groups(self, **kwargs) -> tuple[Sequence[nn.Parameter], Sequence[nn.Parameter]]:
        """Return a tuple of two sequences, one for parameters with weight decay and one for parameters without weight decay.
        Default separation:
        - weight decay: all parameters which have > 1 dimensions.
        - no weight decay: all parameters which have = 1 dimension, e.g. biases.
        """

        decay = set()
        no_decay = set()
        for name, param in self.named_parameters():
            if param.requires_grad:
                # TODO: this is a hack based on module naming. It fixes the issue with the previous heuristic using
                #  parameter dim, which fails for FSDP because it flattens parameters so ndim is always 1.
                if 'weight' in name and 'norm' not in name:
                    decay.add(param)
                else:
                    no_decay.add(param)

        return tuple(decay), tuple(no_decay)
Return a tuple of two sequences, one for parameters with weight decay and one for parameters without weight decay.
        Default separation:
        - weight decay: all parameters which have > 1 dimensions.
        - no weight decay: all parameters which have = 1 dimension, e.g. biases.
        
xlstm/utils.py