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