Term

Contents

Term#

class braintools.fit.Term(fn, weight=1.0, key=None, name=None)#

A weighted loss term of an Objective.

Parameters:

  • fn (Callable[[Any, Any], Any]) – fn(prediction, target) -> scalar, the loss.

  • weight (Any) – Multiplier of the loss.

  • key (Hashable | None) – When given, the term evaluates fn(prediction[key], target[key]), e.g. to pick one output of a dict-valued prediction.

  • name (str | None) – Label used by Objective.breakdown(). Defaults to key when given, otherwise to the name of fn.

Raises:

TypeError – If fn is not callable or weight is not a scalar.

See also

Objective

The sum of several terms.

Fitter

Fits a model to data by minimising an objective.

Notes

A term is a plain callable: it holds no JAX state, so it can be used inside jax.jit, jax.vmap and jax.grad as long as fn can. 2.0 * term returns a new term with the weight doubled; term + other returns an Objective.

Examples

>>> import jax.numpy as jnp
>>> from braintools.fit import Term
>>> mse = lambda p, t: jnp.mean((p - t) ** 2)
>>> term = 2.0 * Term(mse, name='mse')
>>> term
Term(name='mse', weight=2.0, key=None)
>>> print(term(jnp.array([1., 2.]), jnp.array([1., 4.])))
4.0
Contents