Skip to content

gegenbauer_expansion

Bases: transformation

Source code in tinybig/expansion/orthogonal_polynomial_expansion.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
class gegenbauer_expansion(transformation):

    def __init__(self, name='gegenbauer_polynomial_expansion', d: int = 2, alpha: float = 1.0, *args, **kwargs):
        super().__init__(name=name, *args, **kwargs)
        self.d = d
        self.alpha = alpha

    def calculate_D(self, m: int):
        return m * self.d

    def forward(self, x: torch.Tensor, device='cpu', *args, **kwargs):
        b, m = x.shape
        x = self.pre_process(x=x, device=device)

        # base case: order 0
        expansion = torch.ones(size=[x.size(0), x.size(1), self.d + 1]).to(device)
        # base case: order 1
        if self.d > 0:
            expansion[:, :, 1] = 2 * self.alpha * x
        # high-order cases
        for n in range(2, self.d + 1):
            expansion[:, :, n] = (n-1+self.alpha)/n * 2*x * expansion[:, :, n-1].clone() - (n-2+2*self.alpha)/n * expansion[:, :, n-2].clone()

        expansion = expansion[:, :, 1:].contiguous().view(x.size(0), -1)

        assert expansion.shape == (b, self.calculate_D(m=m))
        return self.post_process(x=expansion, device=device)