MCPcopy Create free account
hub / github.com/RosettaCommons/RFdiffusion / truncated_normal

Function truncated_normal

rfdiffusion/util_module.py:11–24  ·  view source on GitHub ↗
(uniform, mu=0.0, sigma=1.0, a=-2, b=2)

Source from the content-addressed store, hash-verified

9
10def init_lecun_normal(module, scale=1.0):
11 def truncated_normal(uniform, mu=0.0, sigma=1.0, a=-2, b=2):
12 normal = torch.distributions.normal.Normal(0, 1)
13
14 alpha = (a - mu) / sigma
15 beta = (b - mu) / sigma
16
17 alpha_normal_cdf = normal.cdf(torch.tensor(alpha))
18 p = alpha_normal_cdf + (normal.cdf(torch.tensor(beta)) - alpha_normal_cdf) * uniform
19
20 v = torch.clamp(2 * p - 1, -1 + 1e-8, 1 - 1e-8)
21 x = mu + sigma * np.sqrt(2) * torch.erfinv(v)
22 x = torch.clamp(x, a, b)
23
24 return x
25
26 def sample_truncated_normal(shape, scale=1.0):
27 stddev = np.sqrt(scale/shape[-1])/.87962566103423978 # shape[-1] = fan_in

Callers 1

sample_truncated_normalFunction · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected