r""" Args: x: `(*, B, C)` Returns: y: `(*, B, C)`
(self, x: torch.Tensor)
| 39 | self.rng = np.random.default_rng() |
| 40 | |
| 41 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 42 | r""" |
| 43 | Args: |
| 44 | x: |
| 45 | `(*, B, C)` |
| 46 | |
| 47 | Returns: |
| 48 | y: |
| 49 | `(*, B, C)` |
| 50 | """ |
| 51 | if not self.training: |
| 52 | return x # no masking during inference |
| 53 | |
| 54 | ori_x_shape = x.shape |
| 55 | if len(x.shape) < 2: |
| 56 | x = x.unsqueeze(0) # (1, C) |
| 57 | batch_size = x.size(-2) |
| 58 | dim = x.size(-1) |
| 59 | x = x.reshape(-1, batch_size, dim) |
| 60 | |
| 61 | # create mask |
| 62 | chosen_idxs = self.rng.choice( |
| 63 | np.arange(dim), |
| 64 | size=[batch_size], |
| 65 | replace=True, |
| 66 | p=self.probs, |
| 67 | ) # (batch,) int |
| 68 | |
| 69 | idxs = torch.arange(0, dim, device=x.device) # (dim,) |
| 70 | mask = idxs <= (torch.from_numpy(chosen_idxs).unsqueeze(1).to(device=idxs.device)) # (batch, dim) |
| 71 | |
| 72 | y = x * mask.unsqueeze(0) # (-1, batch, dim) |
| 73 | return y.reshape(*ori_x_shape) |