Construct nested dropout layer, which drops the last dimension. Note that it creates a mask of shape (B, C), so if the input x has a dimension larger than 2, the first dimensions will share the same mask, and different instances in the batch uses different masks.
(self, probs: T.Sequence[float])
| 18 | """ |
| 19 | |
| 20 | def __init__(self, probs: T.Sequence[float]): |
| 21 | """ |
| 22 | Construct nested dropout layer, which drops the last dimension. |
| 23 | Note that it creates a mask of shape (B, C), so if the input x |
| 24 | has a dimension larger than 2, the first dimensions will share the |
| 25 | same mask, and different instances in the batch uses different masks. |
| 26 | |
| 27 | Args: |
| 28 | probs: |
| 29 | the probablity of the index to be chosen. If None, uniform probability. |
| 30 | |
| 31 | Input: |
| 32 | x: (*, B, C) |
| 33 | |
| 34 | Output: |
| 35 | y: (*, B, C) |
| 36 | """ |
| 37 | super().__init__() |
| 38 | self.probs = probs |
| 39 | self.rng = np.random.default_rng() |
| 40 | |
| 41 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 42 | r""" |
nothing calls this directly
no outgoing calls
no test coverage detected