def isin(element, test_elements, assume_unique=False, invert=False): element = _remove_jaxarray(element) test_elements = _remove_jaxarray(test_elements) return JaxArray(jnp.isin(element, test_elements, assume_unique, invert))
def recall(top_ids, ground_truth, r_at=20):
num_valid_ground_truth = jnp.sum(ground_truth > -1)
divide_by = jnp.minimum(r_at, num_valid_ground_truth)
sum_recall = jnp.sum(jnp.isin(top_ids[:r_at], ground_truth))
return jnp.where(divide_by == 0, 0.0, sum_recall / divide_by)
def any_in(prediction, target): """For each row in a and b, checks if any element of a is in b.""" return jnp.isin(prediction, target)