neuromancer.slim.butterfly.butterfly module

class neuromancer.slim.butterfly.butterfly.Butterfly(in_size, out_size, bias=True, complex=False, tied_weight=True, increasing_stride=True, ortho_init=False)

Bases: Module

Product of log N butterfly factors, each is a block 2x2 of diagonal matrices. Compatible with torch.nn.Linear.

Parameters:

• in_size – size of input

• out_size – size of output

• bias – If set to False, the layer will not learn an additive bias. Default: True

• complex – whether complex or real

• tied_weight –

  whether the weights in the butterfly factors are tied.

  If True, will have 4N parameters, else will have 2 N log N parameters (not counting bias)

  increasing_stride: whether to multiply with increasing stride (e.g. 2, 4, …, n/2) or

  decreasing stride (e.g., n/2, n/4, …, 2). Note that this only changes the order of multiplication, not how twiddle is stored. In other words, twiddle[@log_stride] always stores the twiddle for @stride.

• ortho_init – whether the weight matrix should be initialized to be orthogonal/unitary.

extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(input)

Parameters:

input – (batch, in_size) if real or (batch, in_size, 2) if complex

Returns:

(batch, out_size) if real or (batch, out_size, 2) if complex

Return type:

output

reset_parameters()

Initialize bias the same way as torch.nn.Linear.