Function loss

mbd/blackbox/mbd_mnist.py:158–161

(params, batch)

def loss(params, batch):
    inputs, targets = batch
    preds = predict(params, inputs)
    return -jnp.mean(jnp.sum(preds * targets, axis=1))


def accuracy(params, batch):