(
test_case,
shape,
normalized_shape,
affine=True,
eps=1e-6,
dtype=flow.float32,
device="cuda",
backward=True,
)
| 46 | |
| 47 | |
| 48 | def _test_layer_norm( |
| 49 | test_case, |
| 50 | shape, |
| 51 | normalized_shape, |
| 52 | affine=True, |
| 53 | eps=1e-6, |
| 54 | dtype=flow.float32, |
| 55 | device="cuda", |
| 56 | backward=True, |
| 57 | ): |
| 58 | np_x = np.random.randn(*shape).astype(np.float32) |
| 59 | if affine: |
| 60 | np_weight = np.random.randn(*normalized_shape).astype(np.float32) |
| 61 | np_bias = np.random.randn(*normalized_shape).astype(np.float32) |
| 62 | |
| 63 | # torch process |
| 64 | torch_dtype = torch.float16 if dtype is flow.float16 else torch.float32 |
| 65 | torch_x = torch.tensor(np_x).to(device=device, dtype=torch_dtype) |
| 66 | if backward: |
| 67 | torch_x.requires_grad_(True) |
| 68 | torch_weight = None |
| 69 | torch_bias = None |
| 70 | if affine: |
| 71 | torch_weight = torch.tensor(np_weight).to(device=device, dtype=torch_dtype) |
| 72 | torch_bias = torch.tensor(np_bias).to(device=device, dtype=torch_dtype) |
| 73 | if backward: |
| 74 | torch_weight.requires_grad_(True) |
| 75 | torch_bias.requires_grad_(True) |
| 76 | torch_y = torch.nn.functional.layer_norm( |
| 77 | torch_x, normalized_shape, torch_weight, torch_bias, eps |
| 78 | ) |
| 79 | |
| 80 | if backward: |
| 81 | np_rand_init_grad = np.random.randn(*tuple(torch_y.shape)).astype(np.float32) |
| 82 | torch_rand_init_grad = torch.tensor(np_rand_init_grad).to( |
| 83 | device=device, dtype=torch_dtype |
| 84 | ) |
| 85 | (torch_y * torch_rand_init_grad).sum().backward() |
| 86 | |
| 87 | torch_x_grad = torch_x.grad.detach().cpu().numpy() |
| 88 | if affine: |
| 89 | torch_weight_grad = torch_weight.grad.detach().cpu().numpy() |
| 90 | torch_bias_grad = torch_bias.grad.detach().cpu().numpy() |
| 91 | |
| 92 | torch_y = torch_y.detach().cpu().numpy() |
| 93 | |
| 94 | # oneflow process |
| 95 | x = flow.tensor(np_x).to(device=device, dtype=dtype) |
| 96 | if backward: |
| 97 | x.requires_grad_(True) |
| 98 | weight = None |
| 99 | bias = None |
| 100 | if affine: |
| 101 | weight = flow.tensor(np_weight).to(device=device, dtype=dtype) |
| 102 | bias = flow.tensor(np_bias).to(device=device, dtype=dtype) |
| 103 | if backward: |
| 104 | weight.requires_grad_(True) |
| 105 | bias.requires_grad_(True) |
no test coverage detected