Objective

class braintools.fit.Objective(*terms)

A weighted sum of loss terms.

Parameters:

*terms (Term | Objective | Callable) – The terms. Objectives are flattened into their terms; a bare callable fn(prediction, target) is wrapped as Term(fn).

Raises:

• ValueError – If no term is given.

• TypeError – If a term is neither a Term, an Objective nor a callable.

See also

Term

One weighted loss.

Fitter

Fits a model to data by minimising an objective.

Notes

objective(prediction, target) returns the sum of the weighted terms. The objective is a plain callable that works under jax.jit, jax.vmap and jax.grad whenever its terms do. Objectives support + with terms and objectives and multiplication by a scalar, which scales every weight.

Examples

>>> import jax.numpy as jnp
>>> from braintools.fit import Term
>>> mse = lambda p, t: jnp.mean((p - t) ** 2)
>>> mae = lambda p, t: jnp.mean(jnp.abs(p - t))
>>> obj = 0.5 * Term(mse, key='v', name='mse_v') + Term(mae, key='r')
>>> pred = {'v': jnp.array([0., 2.]), 'r': jnp.array([1., 1.])}
>>> target = {'v': jnp.array([0., 0.]), 'r': jnp.array([0., 2.])}
>>> print(obj(pred, target))
2.0
>>> {k: float(v) for k, v in obj.breakdown(pred, target).items()}
{'mse_v': 1.0, 'r': 1.0}

breakdown(prediction, target)

Return the weighted value of every term, for logging.

Parameters:

• prediction (Any) – As for __call__().

• target (Any) – As for __call__().

Returns:

{name: weight * fn(...)} in term order, keyed by the unique labels of names (repeated names get the first unused suffix _1, _2, …); the values sum to the objective.

Return type:

Dict[str, Any]

property names: Tuple[str, ...]

Unique labels of the terms, as used by breakdown().

The first term with a given name keeps it; every later term with the same name gets the first suffix _1, _2, … that no other label uses, so the labels never collide.