Minor syntactic change.
[pytorch] / poly.py
1 #!/usr/bin/env python
2
3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
5
6 # Written by Francois Fleuret <francois@fleuret.org>
7
8 ######################################################################
9
10 import torch
11
12 def pol_prod(a, b):
13     m = a[:, None] * b[None, :]
14     mm = m.new()
15     mm.set_(m.storage(), 0, (m.size(0), m.size(0) + m.size(1) - 1), (m.size(1) - 1, 1))
16     k = torch.arange(a.size(0))[:, None] + torch.arange(b.size(0))[None, :]
17     kk = k.new()
18     kk.set_(k.storage(), 0, (k.size(0), k.size(0) + k.size(1) - 1), (k.size(1) - 1, 1))
19     q = (kk == torch.arange(a.size(0) + b.size(0) - 1)[None, :])
20     return (mm * q).sum(0)
21
22 def pol_eval(a, x):
23     d = torch.arange(a.size(0))
24     return (x[:, None].pow(d[None, :]) * a[None, :]).sum(1)
25
26 def pol_prim(a):
27     n = torch.arange(a.size(0) + 1).float()
28     n[1:] = a / n[1:]
29     return n
30
31 ######################################################################
32
33 if __name__ == '__main__':
34     a = torch.tensor([1., 2., 3.])
35     b = torch.tensor([2., 5.])
36     print(pol_prod(a, b))
37     print(pol_prim(b))
38     print(pol_eval(a, torch.tensor([0.0, 1.0, 2.0])))