If trunc is None, solve `min sum_i w_i * |a * x_i - y_i|`, otherwise solve `min sum_i min(trunc, w_i * |a * x_i - y_i|)`. w_i must be >= 0. ### Parameters: - `x`: tensor of shape (..., n) - `y`: tensor of shape (..., n) - `w`: tensor of shape (..., n) - `trunc`: op
(x: torch.Tensor, y: torch.Tensor, w: torch.Tensor, trunc: Optional[Union[float, torch.Tensor]] = None, eps: float = 1e-7)
| 50 | |
| 51 | |
| 52 | def align(x: torch.Tensor, y: torch.Tensor, w: torch.Tensor, trunc: Optional[Union[float, torch.Tensor]] = None, eps: float = 1e-7) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: |
| 53 | """ |
| 54 | If trunc is None, solve `min sum_i w_i * |a * x_i - y_i|`, otherwise solve `min sum_i min(trunc, w_i * |a * x_i - y_i|)`. |
| 55 | |
| 56 | w_i must be >= 0. |
| 57 | |
| 58 | ### Parameters: |
| 59 | - `x`: tensor of shape (..., n) |
| 60 | - `y`: tensor of shape (..., n) |
| 61 | - `w`: tensor of shape (..., n) |
| 62 | - `trunc`: optional, float or tensor of shape (..., n) or None |
| 63 | |
| 64 | ### Returns: |
| 65 | - `a`: tensor of shape (...), differentiable |
| 66 | - `loss`: tensor of shape (...), value of loss function at `a`, detached |
| 67 | - `index`: tensor of shape (...), where a = y[idx] / x[idx] |
| 68 | """ |
| 69 | if trunc is None: |
| 70 | x, y, w = torch.broadcast_tensors(x, y, w) |
| 71 | sign = torch.sign(x) |
| 72 | x, y = x * sign, y * sign |
| 73 | y_div_x = y / x.clamp_min(eps) |
| 74 | y_div_x, argsort = y_div_x.sort(dim=-1) |
| 75 | |
| 76 | wx = torch.gather(x * w, dim=-1, index=argsort) |
| 77 | derivatives = 2 * wx.cumsum(dim=-1) - wx.sum(dim=-1, keepdim=True) |
| 78 | search = torch.searchsorted(derivatives, torch.zeros_like(derivatives[..., :1]), side='left').clamp_max(derivatives.shape[-1] - 1) |
| 79 | |
| 80 | a = y_div_x.gather(dim=-1, index=search).squeeze(-1) |
| 81 | index = argsort.gather(dim=-1, index=search).squeeze(-1) |
| 82 | loss = (w * (a[..., None] * x - y).abs()).sum(dim=-1) |
| 83 | |
| 84 | else: |
| 85 | # Reshape to (batch_size, n) for simplicity |
| 86 | x, y, w = torch.broadcast_tensors(x, y, w) |
| 87 | batch_shape = x.shape[:-1] |
| 88 | batch_size = math.prod(batch_shape) |
| 89 | x, y, w = x.reshape(-1, x.shape[-1]), y.reshape(-1, y.shape[-1]), w.reshape(-1, w.shape[-1]) |
| 90 | |
| 91 | sign = torch.sign(x) |
| 92 | x, y = x * sign, y * sign |
| 93 | wx, wy = w * x, w * y |
| 94 | xyw = torch.stack([x, y, w], dim=-1) # Stacked for convenient gathering |
| 95 | |
| 96 | y_div_x = A = y / x.clamp_min(eps) |
| 97 | B = (wy - trunc) / wx.clamp_min(eps) |
| 98 | C = (wy + trunc) / wx.clamp_min(eps) |
| 99 | with torch.no_grad(): |
| 100 | # Caculate prefix sum by orders of A, B, C |
| 101 | A, A_argsort = A.sort(dim=-1) |
| 102 | Q_A = torch.cumsum(torch.gather(wx, dim=-1, index=A_argsort), dim=-1) |
| 103 | A, Q_A = _pad_inf(A), _pad_cumsum(Q_A) # Pad [-inf, A1, ..., An, inf] and [0, Q1, ..., Qn, Qn] to handle edge cases. |
| 104 | |
| 105 | B, B_argsort = B.sort(dim=-1) |
| 106 | Q_B = torch.cumsum(torch.gather(wx, dim=-1, index=B_argsort), dim=-1) |
| 107 | B, Q_B = _pad_inf(B), _pad_cumsum(Q_B) |
| 108 | |
| 109 | C, C_argsort = C.sort(dim=-1) |
no test coverage detected