Compute the cumulative sum of w, assuming all weight vectors sum to 1. The output's size on the last dimension is one greater than that of the input, because we're computing the integral corresponding to the endpoints of a step function, not the integral of the interior/bin values. Args:
(w)
| 129 | |
| 130 | |
| 131 | def integrate_weights_np(w): |
| 132 | """Compute the cumulative sum of w, assuming all weight vectors sum to 1. |
| 133 | |
| 134 | The output's size on the last dimension is one greater than that of the input, |
| 135 | because we're computing the integral corresponding to the endpoints of a step |
| 136 | function, not the integral of the interior/bin values. |
| 137 | |
| 138 | Args: |
| 139 | w: Tensor, which will be integrated along the last axis. This is assumed to |
| 140 | sum to 1 along the last axis, and this function will (silently) break if |
| 141 | that is not the case. |
| 142 | |
| 143 | Returns: |
| 144 | cw0: Tensor, the integral of w, where cw0[..., 0] = 0 and cw0[..., -1] = 1 |
| 145 | """ |
| 146 | cw = np.minimum(1, np.cumsum(w[..., :-1], axis=-1)) |
| 147 | shape = cw.shape[:-1] + (1,) |
| 148 | # Ensure that the CDF starts with exactly 0 and ends with exactly 1. |
| 149 | cw0 = np.concatenate([np.zeros(shape), cw, |
| 150 | np.ones(shape)], axis=-1) |
| 151 | return cw0 |
| 152 | |
| 153 | |
| 154 | def invert_cdf(u, t, w_logits): |