neuromancer.slim.butterfly package

Submodules

neuromancer.slim.butterfly.benchmark module

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.

neuromancer.slim.butterfly.butterfly_multiply module

class neuromancer.slim.butterfly.butterfly_multiply.ButterflyFactorMult(*args, **kwargs)

Bases: Function

static backward(ctx, grad)

Parameters:

grad – (batch_size, 2, n) if real or (batch_size, 2, n, 2) if complex

Returns:

(2, 2, n) if real or (2, 2, n, 2) if complex d_input: (batch_size, 2, n) if real or (batch_size, 2, n, 2) if complex

Return type:

d_twiddle

static forward(ctx, twiddle, input)

Multiply by a single factor. :param twiddle: (2, 2, n) if real or (2, 2, n, 2) if complex :param input: (batch_size, 2, n) if real or (batch_size, 2, n, 2) if complex

Returns:

(batch_size, 2, n) if real or (batch_size, 2, n, 2) if complex

Return type:

output

class neuromancer.slim.butterfly.butterfly_multiply.ButterflyMult(*args, **kwargs)

Bases: Function

static backward(ctx, grad)

Parameters:

• grad – (batch_size, nstack, n) if real or (batch_size, nstack, n, 2) if complex

• twiddle – (nstack, n - 1, 2, 2) if real or (nstack, n - 1, 2, 2, 2) if complex

• backward (output + intermediate values for) – (log n + 1, batch_size, nstack, n) if real or (log n + 1, batch_size, nstack, n, 2) if complex

Returns:

(nstack, n - 1, 2, 2) if real or (nstack, n - 1, 2, 2, 2) if complex d_input: (batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

d_twiddle

static forward(ctx, twiddle, input, increasing_stride=True)

Parameters:

• twiddle – (nstack, n - 1, 2, 2) if real or (nstack, n - 1, 2, 2, 2) if complex

• input – (batch_size, n) if real or (batch_size, n, 2) if complex

• 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.

Returns:

(batch_size, nstack, n) if real or (batch_size, nstack, n, 2) if complex

Return type:

output

class neuromancer.slim.butterfly.butterfly_multiply.ButterflyMultInplace(*args, **kwargs)

Bases: Function

static backward(ctx, grad)

Defines a formula for differentiating the operation with backward mode automatic differentiation (alias to the vjp function).

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static forward(ctx, twiddle, input, increasing_stride=True)

Experimental in-place implementation that does not store intermediate results. Instead, the intermediate results are computed from the output during the backward pass. :param twiddle: (n - 1, 2, 2) if real or (n - 1, 2, 2, 2) if complex :param input: (batch_size, n) if real or (batch_size, n, 2) if complex :param 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.

Returns:

(batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

output

class neuromancer.slim.butterfly.butterfly_multiply.ButterflyMultUntied(*args, **kwargs)

Bases: Function

static backward(ctx, grad)

Parameters:

• grad – (batch_size, nstack, n) if real or (batch_size, nstack, n, 2) if complex

• twiddle – (nstack, log 2, n / 2, 2, 2) if real or (nstack, log 2, n / 2, 2, 2, 2) if complex

• backward (output + intermediate values for) – (log n + 1, batch_size, nstack, n) if real or (log n + 1, batch_size, nstack, n, 2) if complex

Returns:

(nstack, log 2, n / 2, 2, 2) if real or (nstack, log 2, n / 2, 2, 2, 2) if complex d_input: (batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

d_twiddle

static forward(ctx, twiddle, input, increasing_stride=True)

Parameters:

• twiddle – (nstack, log 2, n / 2, 2, 2) if real or (nstack, log 2, n / 2, 2, 2, 2) if complex

• input – (batch_size, n) if real or (batch_size, n, 2) if complex

• 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.

Returns:

(batch_size, nstack, n) if real or (batch_size, nstack, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.butterfly_multiply.butterfly_mult(twiddle, input, increasing_stride=True, return_intermediates=False)

Parameters:

• twiddle – (nstack, n - 1, 2, 2) if real or (nstack, n - 1, 2, 2, 2) if complex

• input – (batch_size, n) if real or (batch_size, n, 2) if complex

• 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.

• return_intermediates – whether to return all the intermediate values computed, for debugging

Returns:

(batch_size, nstack, n) if real or (batch_size, nstack, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.butterfly_multiply.butterfly_mult_factors(twiddle, input, increasing_stride=True, return_intermediates=False)

Implementation that have separate kernels for each factor, for debugging. :param twiddle: (n - 1, 2, 2) if real or (n - 1, 2, 2, 2) if complex :param input: (batch_size, n) if real or (batch_size, n, 2) if complex :param 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.

Parameters:

return_intermediates – whether to return all the intermediate values computed, for debugging

Returns:

(batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.butterfly_multiply.butterfly_mult_torch(twiddle, input, increasing_stride=True, return_intermediates=False)

Parameters:

• twiddle – (nstack, n - 1, 2, 2) if real or (nstack, n - 1, 2, 2, 2) if complex

• input – (batch_size, n) if real or (batch_size, n, 2) if complex

• 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.

• return_intermediates – whether to return all the intermediate values computed, for debugging

Returns:

(batch_size, nstack, n) if real or (batch_size, nstack, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.butterfly_multiply.butterfly_mult_untied(twiddle, input, increasing_stride=True, return_intermediates=False)

Parameters:

• twiddle – (nstack, log n, n / 2, 2, 2) if real or (nstack, log n, n / 2, 2, 2, 2) if complex

• input – (batch_size, n) if real or (batch_size, n, 2) if complex

• 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.

• return_intermediates – whether to return all the intermediate values computed, for debugging

Returns:

(batch_size, nstack, n) if real or (batch_size, nstack, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.butterfly_multiply.butterfly_mult_untied_torch(twiddle, input, increasing_stride=True, return_intermediates=False)

Parameters:

• twiddle – (nstack, log n, n / 2, 2, 2) if real or (nstack, log n, n / 2, 2, 2, 2) if complex

• input – (batch_size, n) if real or (batch_size, n, 2) if complex

• 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.

• return_intermediates – whether to return all the intermediate values computed, for debugging

Returns:

(batch_size, nstack, n) if real or (batch_size, nstack, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.complex_utils module

Utility functions for handling complex tensors: conjugate and complex_mul. Pytorch (as of 1.0) does not support complex tensors, so we store them as float tensors where the last dimension is 2 (real and imaginary parts).

class neuromancer.slim.butterfly.complex_utils.ComplexMatmulNp(*args, **kwargs)

Bases: Function

Multiply two complex matrices, in numpy. :param X: (n, m, 2) :param Y: (m, p, 2)

Returns:

(n, p, 2)

Return type:

Z

static backward(ctx, grad)

Defines a formula for differentiating the operation with backward mode automatic differentiation (alias to the vjp function).

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static forward(ctx, X, Y)

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

• It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

• See combining-forward-context for more details

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass

@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

• The forward no longer accepts a ctx argument.

• Instead, you must also override the torch.autograd.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.

• See extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

class neuromancer.slim.butterfly.complex_utils.ComplexMul(*args, **kwargs)

Bases: Function

X and Y are complex64 tensors but stored as float32 tensors, with last dimension = 2.

static backward(ctx, grad)

Defines a formula for differentiating the operation with backward mode automatic differentiation (alias to the vjp function).

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static forward(ctx, X, Y)

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

• It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

• See combining-forward-context for more details

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass

@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

• The forward no longer accepts a ctx argument.

• Instead, you must also override the torch.autograd.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.

• See extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

class neuromancer.slim.butterfly.complex_utils.Conjugate(*args, **kwargs)

Bases: Function

X is a complex64 tensors but stored as float32 tensors, with last dimension = 2.

static backward(ctx, grad)

Defines a formula for differentiating the operation with backward mode automatic differentiation (alias to the vjp function).

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static forward(ctx, X)

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

• It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

• See combining-forward-context for more details

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass

@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

• The forward no longer accepts a ctx argument.

• Instead, you must also override the torch.autograd.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.

• See extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

neuromancer.slim.butterfly.complex_utils.complex_matmul_torch(X, Y)

Multiply two complex matrices. :param X: (…, n, m, 2) :param Y: (…, m, p, 2)

Returns:

(…, n, p, 2)

Return type:

Z

neuromancer.slim.butterfly.complex_utils.complex_mul_numpy(X, Y)

neuromancer.slim.butterfly.complex_utils.complex_mul_torch(X, Y)

neuromancer.slim.butterfly.complex_utils.conjugate_torch(X)

neuromancer.slim.butterfly.complex_utils.cupy2torch(tensor)

neuromancer.slim.butterfly.complex_utils.real_to_complex(X)

A version of X that’s complex (i.e., last dimension is 2). :param X: (…) tensor

Returns:

(…, 2) tensor

Return type:

X_complex

neuromancer.slim.butterfly.complex_utils.test_complex_mm()

neuromancer.slim.butterfly.complex_utils.test_complex_mul()

neuromancer.slim.butterfly.complex_utils.torch2cupy(tensor)

neuromancer.slim.butterfly.complex_utils.torch2numpy(X)

Convert a torch float32 tensor to a numpy array, sharing the same memory.

neuromancer.slim.butterfly.permutation module

class neuromancer.slim.butterfly.permutation.FixedPermutation(permutation)

Bases: Module

forward(input)

Parameters:

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

Returns:

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

Return type:

output

class neuromancer.slim.butterfly.permutation.Permutation(size, share_logit=False, increasing_stride=False)

Bases: Module

Product of log N permutation factors.

Parameters:

• size – size of input (and of output)

• share_logit – whether the logits in the permutation factors are shared. If True, will have 4N parameters, else will have 2 N log N parameters (not counting bias)

• increasing_stride – whether to multiply from smaller stride to larger stride, or in the reverse order.

argmax()

Returns:

(self.size, ) array of int, the most probable permutation.

Return type:

p

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, size) if real or (batch, size, 2) if complex

Returns:

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

Return type:

output

class neuromancer.slim.butterfly.permutation.PermutationFactor(size)

Bases: Module

A single permutation factor.

Parameters:

size – size of input (and of output)

argmax()

Returns:

(self.size, ) array of int, the most probable permutation.

Return type:

p

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, size) if real or (batch, size, 2) if complex

Returns:

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

Return type:

output

neuromancer.slim.butterfly.permutation_multiply module

class neuromancer.slim.butterfly.permutation_multiply.PermutationFactorEvenOddMult(*args, **kwargs)

Bases: Function

static backward(ctx, grad)

Parameters:

grad – (batch_size, n) if real or (batch_size, n, 2) if complex

Returns:

real number d_input: (batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

d_p

static forward(ctx, p, input)

Multiply by a single permutation factor that separates the even and the odd (with weight). :param p: real number between 0.0 and 1.0 :param input: (batch_size, n) if real or (batch_size, n, 2) if complex

Returns:

(batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

output

class neuromancer.slim.butterfly.permutation_multiply.PermutationFactorReverseMult(*args, **kwargs)

Bases: Function

static backward(ctx, grad)

Parameters:

grad – (batch_size, n) if real or (batch_size, n, 2) if complex

Returns:

real number d_input: (batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

d_p

static forward(ctx, p, input)

Multiply by a single permutation factor that reverses the first and second halves (with weights). :param p: (2, ), must be between 0.0 and 1.0 :param input: (batch_size, n) if real or (batch_size, n, 2) if complex

Returns:

(batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.permutation_multiply.permutation_mult(prob, input, increasing_stride=False, return_intermediates=False)

Multiply by permutation factors, parameterized by the probabilities. :param prob: (nsteps, 3), where prob[:, 0] is the probability of separating the even and odd indices,

and prob[:, 1:3] are the probabilities of reversing the 1st and 2nd halves respectively. Note that stride starts at 4, not 2 (as permutations do nothing at stride 2).

Parameters:

• input – (batch_size, n) if real or (batch_size, n, 2) if complex

• increasing_stride – whether to multiply with increasing stride (e.g. 4, 8, …, n/2) or decreasing stride (e.g., n/2, n/4, …, 4). Note that this only changes the order of multiplication, not how prob is stored. In other words, prob[@log_stride - 1] always stores the probability for @stride.

• return_intermediates – whether to return all the intermediate values computed, for debugging

Returns:

(batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.permutation_multiply.permutation_mult_factors(prob, input, increasing_stride=False, return_intermediates=False)

Multiply by permutation factors, parameterized by the probabilities. :param prob: (nsteps, 3), where prob[:, 0] is the probability of separating the even and odd indices,

and prob[:, 1:3] are the probabilities of reversing the 1st and 2nd halves respectively. Note that stride starts at 4, not 2 (as permutations do nothing at stride 2).

Parameters:

• input – (batch_size, n) if real or (batch_size, n, 2) if complex

• increasing_stride – whether to multiply with increasing stride (e.g. 4, 8, …, n/2) or decreasing stride (e.g., n/2, n/4, …, 4). Note that this only changes the order of multiplication, not how prob is stored. In other words, prob[@log_stride - 1] always stores the probability for @stride.

• return_intermediates – whether to return all the intermediate values computed, for debugging

Returns:

(batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.permutation_multiply.permutation_mult_single(prob, input)

Multiply by a single permutation factor. :param prob: (3, ), where prob[0] is the probability of separating the even and odd indices,

and prob[1:3] are the probabilities of reversing the 1st and 2nd halves respectively.

Parameters:

input – (batch_size, n) if real or (batch_size, n, 2) if complex

Returns:

(batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.permutation_multiply.permutation_mult_single_factor(prob, input)

Multiply by a single permutation factor, parameterized by the probabilities. :param prob: (3, ), where prob[0] is the probability of separating the even and odd indices,

and prob[1:3] are the probabilities of reversing the 1st and 2nd halves respectively.

Parameters:

input – (batch_size, n) if real or (batch_size, n, 2) if complex

Returns:

(batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.permutation_multiply.permutation_mult_single_factor_torch(prob, input)

Multiply by a single permutation factor. :param prob: (3, ), where prob[0] is the probability of separating the even and odd indices,

and prob[1:3] are the probabilities of reversing the 1st and 2nd halves respectively.

Parameters:

input – (batch_size, n) if real or (batch_size, n, 2) if complex

Returns:

(batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.permutation_multiply.permutation_mult_torch(prob, input, increasing_stride=False, return_intermediates=False)

Multiply by permutation factors, parameterized by the probabilities. :param prob: (nsteps, 3), where prob[:, 0] is the probability of separating the even and odd indices,

and prob[:, 1:3] are the probabilities of reversing the 1st and 2nd halves respectively. Note that stride starts at 4, not 2 (as permutations do nothing at stride 2).

Parameters:

• input – (batch_size, n) if real or (batch_size, n, 2) if complex

• increasing_stride – whether to multiply with increasing stride (e.g. 4, 8, …, n/2) or decreasing stride (e.g., n/2, n/4, …, 4). Note that this only changes the order of multiplication, not how prob is stored. In other words, prob[@log_stride - 1] always stores the probability for @stride.

• return_intermediates – whether to return all the intermediate values computed, for debugging

Returns:

(batch_size, n) if real or (batch_size, n, 2) if complex

Return type:

output

neuromancer.slim.butterfly.utils module

neuromancer.slim.butterfly.utils.bitreversal_permutation(n)

Return the bit reversal permutation used in FFT. Parameter:

n: integer, must be a power of 2.

Returns:

bit reversal permutation, numpy array of size n

Return type:

perm

Module contents