Implementation of clamp_with_pushback function.
| 104 | |
| 105 | |
| 106 | class ClampWithPushback(autograd.Function): |
| 107 | """Implementation of clamp_with_pushback function.""" |
| 108 | |
| 109 | @staticmethod |
| 110 | def forward( |
| 111 | ctx: Any, |
| 112 | tensor: torch.Tensor, |
| 113 | min: float | None, |
| 114 | max: float | None, |
| 115 | pushback: float, |
| 116 | ) -> torch.Tensor: |
| 117 | """Apply clamp.""" |
| 118 | if min is not None and max is not None and min >= max: |
| 119 | raise ValueError("Only min < max is supported.") |
| 120 | |
| 121 | ctx.save_for_backward(tensor) |
| 122 | ctx.min = min |
| 123 | ctx.max = max |
| 124 | ctx.pushback = pushback |
| 125 | return torch.clamp(tensor, min=min, max=max) |
| 126 | |
| 127 | @staticmethod |
| 128 | def backward( # type: ignore[override] # Deal with buggy torch annotations. |
| 129 | ctx: Any, grad_in: torch.Tensor |
| 130 | ) -> tuple[torch.Tensor, None, None, None]: |
| 131 | """Compute gradient of clamp with pushback.""" |
| 132 | grad_out = grad_in.clone() |
| 133 | (tensor,) = ctx.saved_tensors |
| 134 | |
| 135 | if ctx.min is not None: |
| 136 | mask_min = tensor < ctx.min |
| 137 | grad_out[mask_min] = -ctx.pushback |
| 138 | |
| 139 | if ctx.max is not None: |
| 140 | mask_max = tensor > ctx.max |
| 141 | grad_out[mask_max] = ctx.pushback |
| 142 | |
| 143 | return grad_out, None, None, None |
| 144 | |
| 145 | |
| 146 | def clamp_with_pushback( |
nothing calls this directly
no outgoing calls
no test coverage detected