def iterate_layers()

in opacus/layers/dp_rnn.py [0:0]

• 9 lines of code
• 4 McCabe index (conditional complexity)

    def iterate_layers(self, *args):
        """
        Iterate through all the layers and through all directions within each layer.

        Arguments should be list-like of length ``num_layers * num_directions`` where
        each element corresponds to (layer, direction) pair. The corresponding elements
        of each of these lists will be iterated over.

        Example:
            num_layers = 3
            bidirectional = True

            for layer, directions in self.iterate_layers(self.cell, h):
                for dir, (cell, hi) in directions:
                    print(layer, dir, hi)

            # 0 0 h[0]
            # 0 1 h[1]
            # 1 0 h[2]
            # 1 1 h[3]
            # 2 0 h[4]
            # 2 1 h[5]

        """
        for layer in range(self.num_layers):
            yield layer, (
                (
                    direction,
                    tuple(arg[self.num_directions * layer + direction] for arg in args),
                )
                for direction in range(self.num_directions)
            )