(i: int, solution: Tensor)
| 241 | del predicate_satisfied |
| 242 | |
| 243 | def loop_body(i: int, solution: Tensor) -> Tensor: |
| 244 | # Loop over the non-sign bits. |
| 245 | bit = jnp.int32(1 << 30 - i) |
| 246 | # pylint: disable-next=unsupported-binary-operation |
| 247 | predicate_satisfied = predicate(solution | bit) |
| 248 | solution = solution | jnp.where(predicate_satisfied, jnp.int32(0), bit) |
| 249 | return solution |
| 250 | |
| 251 | return jax.lax.fori_loop(0, 31, loop_body, solution) |
| 252 |
nothing calls this directly
no test coverage detected