Kernels
iamwyldecat commited on
Commit
552d415
·
1 Parent(s): f5a7d38

chore(poly-norm): remove unnecessary file

Browse files
Files changed (1) hide show
  1. tests/test.py +0 -38
tests/test.py DELETED
@@ -1,38 +0,0 @@
1
- import activation
2
- import torch
3
-
4
-
5
- def norm(x, eps: float) -> torch.Tensor:
6
- return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + eps)
7
-
8
-
9
- def poly_norm(
10
- x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, eps: float
11
- ) -> torch.Tensor:
12
- x = x.float()
13
- return (
14
- weight[0] * norm(x**3, eps)
15
- + weight[1] * norm(x**2, eps)
16
- + weight[2] * norm(x, eps)
17
- + bias
18
- ).to(weight.dtype)
19
-
20
-
21
- dtype = torch.bfloat16
22
- torch.set_default_device("cuda:0")
23
- a = torch.randn(3, 3, dtype=dtype, requires_grad=True)
24
- w = torch.randn(3, dtype=dtype, requires_grad=True)
25
- b = torch.randn(1, dtype=dtype, requires_grad=True)
26
-
27
- a.retain_grad()
28
- w.retain_grad()
29
- b.retain_grad()
30
-
31
- out = activation.poly_norm(a, w, b, 1e-6)
32
- # out = poly_norm(a, w, b, 1e-6)
33
-
34
- out.backward(torch.ones_like(out))
35
-
36
- print(a.grad)
37
- print(w.grad)
38
- print(b.grad)