def f(model_output, targets, weights): # pylint: disable=invalid-name
predictions = jnp.argmax(model_output, axis=-1)
shapes.assert_same_shape(predictions, targets)
position_is_padding = jnp.equal(weights, 0)
position_is_accurate = jnp.logical_or(jnp.equal(predictions, targets),
position_is_padding)
sequence_is_accurate = jnp.all(position_is_accurate, axis=-1)
return jnp.average(sequence_is_accurate)
def _convert_to_nans(x, y):
# if all values in y are non-zeros, return x; otherwise return 0s
return jnp.where(jnp.all(y, keepdims=False), x, x / 0.), y
