(float32_query: Tensor, top_p: Union[float, Tensor])
| 104 | reduce_axis = reducible_probs.ndim - 2 |
| 105 | |
| 106 | def predicate(float32_query: Tensor, top_p: Union[float, Tensor]) -> Tensor: |
| 107 | float32_query = jnp.expand_dims(float32_query, reduce_axis) |
| 108 | # [..., 1, float32_query.shape[-1]] |
| 109 | probability_mass = jnp.sum( |
| 110 | jnp.where(reducible_probs >= float32_query, reducible_probs, 0.0), |
| 111 | axis=reduce_axis, |
| 112 | ) |
| 113 | if not isinstance(top_p, float): |
| 114 | top_p = top_p.reshape((top_p.shape[0], *(1,) * (probability_mass.ndim - 1))) |
| 115 | return probability_mass < top_p |
| 116 | |
| 117 | batched_shape = logits.shape[:-1] # All but the last axis are batched. |
| 118 | threshold = _float32_binary_search(batched_shape, predicate=partial(predicate, top_p=p)) |
no outgoing calls
no test coverage detected