(self, logits_input, means_input, log_scales_input, a=torch.as_tensor(-1.0 + eps, dtype=torch.float32), b=torch.as_tensor(1.0 - eps, dtype=torch.float32), interval=0.05)
| 76 | return means + scales * torch.erfinv(2 * p - 1) * math.sqrt(2) |
| 77 | |
| 78 | def sample(self, logits_input, means_input, log_scales_input, a=torch.as_tensor(-1.0 + eps, dtype=torch.float32), b=torch.as_tensor(1.0 - eps, dtype=torch.float32), interval=0.05): |
| 79 | |
| 80 | means, log_scales, probs = self.get_mix_params(logits_input, means_input, log_scales_input) # [B, N*n_sample, dim], [B, N*n_sample, 1] |
| 81 | |
| 82 | if not self.training or not self.perform_sampling: |
| 83 | probs_samples = self.log_pdf_fn(means, means, log_scales).exp().mean(dim=-1, keepdim=True) |
| 84 | probs = probs * probs_samples.tanh() |
| 85 | return means, probs |
| 86 | |
| 87 | if a is None: |
| 88 | cdf_lower = torch.zeros_like(means, dtype=torch.float32) |
| 89 | else: |
| 90 | a = a.to(means.device) |
| 91 | a = a.expand_as(means) |
| 92 | cdf_lower = self.cdf_fn(a, means, log_scales) |
| 93 | |
| 94 | if b is None: |
| 95 | cdf_upper = torch.ones_like(means, dtype=torch.float32) |
| 96 | else: |
| 97 | b = b.to(means.device) |
| 98 | b = b.expand_as(means) |
| 99 | cdf_upper = self.cdf_fn(b, means, log_scales) |
| 100 | |
| 101 | # sampling |
| 102 | cdf_delta = cdf_upper - cdf_lower |
| 103 | cdf_mid = (cdf_upper + cdf_lower) / 2 |
| 104 | |
| 105 | u_normal = torch.rand_like(means).clamp(eps, 1.0 - eps) |
| 106 | u_normal = u_normal * cdf_delta + cdf_lower # normal condition |
| 107 | u_edge = 0.5 # dummy value |
| 108 | u = torch.where(cdf_delta > eps, u_normal, u_edge) |
| 109 | u = torch.where(log_scales > self.log_scales_min, u, u_edge) |
| 110 | val = self.icdf_fn(u, means, log_scales) |
| 111 | val = torch.where(cdf_delta > eps, val, a) |
| 112 | |
| 113 | select_val = means |
| 114 | if b is not None: |
| 115 | select_val = torch.where(select_val < b, select_val, b) |
| 116 | if a is not None: |
| 117 | select_val = torch.where(select_val > a, select_val, a) |
| 118 | val = torch.where(log_scales > self.log_scales_min, val, select_val) |
| 119 | |
| 120 | # cdf interval |
| 121 | left = val - interval |
| 122 | if a is not None: |
| 123 | left = torch.max(left, a) |
| 124 | right = val + interval |
| 125 | if b is not None: |
| 126 | right = torch.min(right, b) |
| 127 | cdf_left = self.cdf_fn(left, means, log_scales) |
| 128 | cdf_right = self.cdf_fn(right, means, log_scales) |
| 129 | cdf = (cdf_right - cdf_left) / (cdf_delta + 1e-5) |
| 130 | probs = probs * cdf.min(dim=-1).values.unsqueeze(-1) |
| 131 | |
| 132 | return val, probs |
| 133 |
no test coverage detected