| 657 | |
| 658 | |
| 659 | class DropPath(nn.Module): |
| 660 | dropout_prob: float = 0.0 |
| 661 | deterministic: Optional[bool] = None |
| 662 | |
| 663 | @nn.compact |
| 664 | def __call__(self, input, deterministic=None): |
| 665 | deterministic = nn.merge_param( |
| 666 | "deterministic", self.deterministic, deterministic |
| 667 | ) |
| 668 | if deterministic: |
| 669 | return input |
| 670 | keep_prob = 1 - self.dropout_prob |
| 671 | shape = (input.shape[0],) + (1,) * (input.ndim - 1) |
| 672 | rng = self.make_rng("drop_path") |
| 673 | random_tensor = keep_prob + jax.random.uniform(rng, shape, dtype=jnp.float32) |
| 674 | random_tensor = jnp.floor(random_tensor) |
| 675 | return jnp.divide(input, keep_prob) * random_tensor |