| 65 | |
| 66 | |
| 67 | def temp_sample(self, orig_pdist, temp): |
| 68 | locs, scales, weights = orig_pdist |
| 69 | if temp is None: |
| 70 | component_samples = locs + scales * T.randn_like(scales) |
| 71 | mixture_samples = F.gumbel_softmax(weights, hard=True) |
| 72 | sampled = (component_samples * mixture_samples.unsqueeze(-1)).sum(dim=-2) |
| 73 | elif isinstance(temp, tuple): |
| 74 | assert len(temp) == 2 |
| 75 | categorical_temp, gaussian_temp = temp |
| 76 | component_samples = locs + scales * gaussian_temp * T.randn_like(scales) |
| 77 | mixture_samples = F.gumbel_softmax(weights / categorical_temp, hard=True) |
| 78 | sampled = (component_samples * mixture_samples.unsqueeze(-1)).sum(dim=-2) |
| 79 | else: |
| 80 | component_samples = locs + scales * temp * T.randn_like(scales) |
| 81 | mixture_samples = F.gumbel_softmax(weights / temp, hard=True) |
| 82 | sampled = (component_samples * mixture_samples.unsqueeze(-1)).sum(dim=-2) |
| 83 | return sampled |
| 84 | |
| 85 | |
| 86 | class GPTOutput(nn.Module): |