Simple cumulative accuracy metric using Haiku and Jax.

haiku-accuracy.py

class Accuracy(hk.Module):
    def __init__(self, y_true: jnp.ndarray, y_pred: jnp.ndarray) -> jnp.ndarray:

        total = hk.get_state(
            "total", shape=[], dtype=jnp.float32, init=hk.initializers.Constant(0)
        )
        count = hk.get_state(
            "count", shape=[], dtype=jnp.int64, init=hk.initializers.Constant(0)
        )

        total += jnp.mean(y_true == jnp.argmax(y_pred, axis=-1))
        count += jnp.prod(y_true.shape)

        hk.set_state("total", total)
        hk.set_state("count", count)

        return total / count