your_name/jax_example

@measurement.evaluate(params=0, batch=1)
def evaluate(params, batch):
    logits = model.apply({"params": params}, batch)
    return jax.nn.softmax(logits, axis=-1).max(axis=-1)

@measurement.aggregate
def aggregate(outputs):
    values = [value for batch in outputs for value in batch]
    return sum(values) / len(values)