def _get_identity(op, dtype):
"""Get an appropriate identity for a given operation in a given dtype."""
if op is lax.scatter_add:
return 0
elif op is lax.scatter_mul:
return 1
elif op is lax.scatter_min:
if jnp.issubdtype(dtype, jnp.integer):
return jnp.iinfo(dtype).max
return float('inf')
elif op is lax.scatter_max:
if jnp.issubdtype(dtype, jnp.integer):
return jnp.iinfo(dtype).min
return -float('inf')
else:
raise ValueError(f"Unrecognized op: {op}")
def loop_body(state):
agenda, blocks, eigenvectors = state
(offset, b), agenda = agenda.pop()
which = jnp.where(buckets < b, jnp.iinfo(jnp.int32).max, buckets)
choice = jnp.argmin(which)
return lax.switch(choice, branches, offset, b, agenda, blocks, eigenvectors)