Confusion matrix

[krzysztofrusek]

Hi

I propose a new endpoint jax.nn.confusion_matrix with the following implementation

@partial(jax.jit, static_argnames=('n'))
def confusion_matrix(labels, predictions, n):
    cm, _ = jax.lax.scan(
        lambda carry, pair: (carry.at[pair].add(1), None), 
        jnp.zeros((n,n), dtype=jnp.uint32), 
        (labels, predictions)
        )
    return cm

and the test:

a = jnp.array([1,2,0,1,2,1,2])
b = jnp.array([1,2,0,1,2,1,0])
k = confusion_matrix(a,b,3)
expected = jnp.array([[1,0,0],[0,3,0],[1,0,2]])

np.allclose(k, expected)

I could prepare a pull request, but I want to discuss the design before.