Robust torch.where function to avoid NaN in backward pass. See https://github.com/pytorch/pytorch/issues/68425 Args: condition: When True (nonzero), yield branch_true_func(input), otherwise yield branch_false_func(input) input: The input tensor for torch.where
(
condition: torch.Tensor,
input: torch.Tensor,
branch_true_func: Callable[[torch.Tensor], torch.Tensor],
branch_false_func: Callable[[torch.Tensor], torch.Tensor],
branch_true_safe_value: float | None = None,
branch_false_safe_value: float | None = None,
)
| 12 | |
| 13 | |
| 14 | def robust_where( |
| 15 | condition: torch.Tensor, |
| 16 | input: torch.Tensor, |
| 17 | branch_true_func: Callable[[torch.Tensor], torch.Tensor], |
| 18 | branch_false_func: Callable[[torch.Tensor], torch.Tensor], |
| 19 | branch_true_safe_value: float | None = None, |
| 20 | branch_false_safe_value: float | None = None, |
| 21 | ) -> torch.Tensor: |
| 22 | """Robust torch.where function to avoid NaN in backward pass. |
| 23 | |
| 24 | See https://github.com/pytorch/pytorch/issues/68425 |
| 25 | |
| 26 | Args: |
| 27 | condition: When True (nonzero), yield branch_true_func(input), |
| 28 | otherwise yield branch_false_func(input) |
| 29 | input: The input tensor for torch.where |
| 30 | branch_true_func: Callable for values at indices where condition is True. |
| 31 | branch_false_func: Callable for values at indices where condition is False. |
| 32 | branch_true_safe_value: Safe value to replace the true branch. |
| 33 | branch_false_safe_value: Safe value to replace the false branch. |
| 34 | """ |
| 35 | input_1 = input |
| 36 | input_2 = input |
| 37 | if branch_true_safe_value is not None: |
| 38 | input_1 = torch.where(condition, input_1, branch_true_safe_value) |
| 39 | if branch_false_safe_value is not None: |
| 40 | input_2 = torch.where(~condition, input_2, branch_false_safe_value) |
| 41 | return torch.where( |
| 42 | condition, |
| 43 | branch_true_func(input_1), |
| 44 | branch_false_func(input_2), |
| 45 | ) |
no test coverage detected