probly.calibration.scaling.flax_base

Base Implementation For Platt-, Vector- and Temperature Scaling.

Classes

ScalerFlax(*args, **kwargs)

Base class for Flax Scaling Implementations.
class probly.calibration.scaling.flax_base.ScalerFlax(*args, **kwargs)[source]

Bases: Module, ABC

Base class for Flax Scaling Implementations.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

Any

__call__(x)[source]

Call the scaler and returne scaled logits.

Parameters:

x (Array) – The input the scaled logits are produced from

Returns:

The scaled logits based on the input.

Return type:

scaled_logits

eval(**attributes)

Sets the Module to evaluation mode.

eval uses set_attributes to recursively set attributes deterministic=True and use_running_average=True of all nested Modules that have these attributes. Its primarily used to control the runtime behavior of the Dropout and BatchNorm Modules.

Example:

>>> from flax import nnx
...
>>> class Block(nnx.Module):
...   def __init__(self, din, dout, *, rngs: nnx.Rngs):
...     self.linear = nnx.Linear(din, dout, rngs=rngs)
...     self.dropout = nnx.Dropout(0.5)
...     self.batch_norm = nnx.BatchNorm(10, rngs=rngs)
...
>>> block = Block(2, 5, rngs=nnx.Rngs(0))
>>> block.dropout.deterministic, block.batch_norm.use_running_average
(False, False)
>>> block.eval()
>>> block.dropout.deterministic, block.batch_norm.use_running_average
(True, True)
Parameters:

**attributes – additional attributes passed to set_attributes.

fit(calibration_set, learning_rate=0.01, max_iter=50)[source]

Optimizes the parameters based on the given calibration set.

Parameters:

  • calibration_set (Iterable[tuple[Array, Array]]) – The data set used for optimizing the parameters

  • learning_rate (float) – The learning rate used by the optimizer

  • max_iter (int) – The maximum steps the optimizer takes

Return type:

None

iter_children()

Iterates over all children Module’s of the current Module. This method is similar to iter_modules(), except it only iterates over the immediate children, and does not recurse further down.

iter_children creates a generator that yields the key and the Module instance, where the key is a string representing the attribute name of the Module to access the corresponding child Module.

Example:

>>> from flax import nnx
...
>>> class SubModule(nnx.Module):
...   def __init__(self, din, dout, rngs):
...     self.linear1 = nnx.Linear(din, dout, rngs=rngs)
...     self.linear2 = nnx.Linear(din, dout, rngs=rngs)
...
>>> class Block(nnx.Module):
...   def __init__(self, din, dout, *, rngs: nnx.Rngs):
...     self.linear = nnx.Linear(din, dout, rngs=rngs)
...     self.submodule = SubModule(din, dout, rngs=rngs)
...     self.dropout = nnx.Dropout(0.5)
...     self.batch_norm = nnx.BatchNorm(10, rngs=rngs)
...
>>> model = Block(2, 5, rngs=nnx.Rngs(0))
>>> for path, module in model.iter_children():
...  print(path, type(module).__name__)
...
batch_norm BatchNorm
dropout Dropout
linear Linear
submodule SubModule
Return type:

Iterator[tuple[Key, Module]]

iter_modules()

Recursively iterates over all nested Module’s of the current Module, including the current Module.

iter_modules creates a generator that yields the path and the Module instance, where the path is a tuple of strings or integers representing the path to the Module from the root Module.

Example:

>>> from flax import nnx
...
>>> class SubModule(nnx.Module):
...   def __init__(self, din, dout, rngs):
...     self.linear1 = nnx.Linear(din, dout, rngs=rngs)
...     self.linear2 = nnx.Linear(din, dout, rngs=rngs)
...
>>> class Block(nnx.Module):
...   def __init__(self, din, dout, *, rngs: nnx.Rngs):
...     self.linear = nnx.Linear(din, dout, rngs=rngs)
...     self.submodule = SubModule(din, dout, rngs=rngs)
...     self.dropout = nnx.Dropout(0.5)
...     self.batch_norm = nnx.BatchNorm(10, rngs=rngs)
...
>>> model = Block(2, 5, rngs=nnx.Rngs(0))
>>> for path, module in model.iter_modules():
...   print(path, type(module).__name__)
...
('batch_norm',) BatchNorm
('dropout',) Dropout
('linear',) Linear
('submodule', 'linear1') Linear
('submodule', 'linear2') Linear
('submodule',) SubModule
() Block
Return type:

Iterator[tuple[tuple[Key, …], Module]]

perturb(name, value, variable_type=<class 'flax.nnx.variablelib.Perturbation'>)

Add an zero-value variable (“perturbation”) to the intermediate value.

The gradient of value would be the same as the gradient of this perturbation variable. Therefore, if you define your loss function with both params and perturbations as standalone arguments, you can get the intermediate gradients of value by running jax.grad on the perturbation variable.

Since the shape of the perturbation value depends on the shape of the input, a perturbation variable is only created after you run a sample input through the model once.

Note

This creates extra dummy variables of the same size as value, thus occupies more memory. Use it only to debug gradients in training.

Example usage:

>>> from flax import nnx
>>> import jax.numpy as jnp

>>> class Model(nnx.Module):
...   def __init__(self, rngs):
...     self.linear1 = nnx.Linear(2, 3, rngs=rngs)
...     self.linear2 = nnx.Linear(3, 4, rngs=rngs)
...   def __call__(self, x):
...     x = self.linear1(x)
...     x = self.perturb('xgrad', x)
...     x = self.linear2(x)
...     return x

>>> x = jnp.ones((1, 2))
>>> y = jnp.ones((1, 4))
>>> model = Model(rngs=nnx.Rngs(0))
>>> assert not hasattr(model, 'xgrad')  # perturbation requires a sample input run
>>> _ = model(x)
>>> assert model.xgrad.value.shape == (1, 3)   # same as the intermediate value
>>> graphdef, params, perturbations = nnx.split(model, nnx.Param, nnx.Perturbation)

>>> # Take gradients on the Param and Perturbation variables
>>> @nnx.grad(argnums=(0, 1))
... def grad_loss(params, perturbations, inputs, targets):
...   model = nnx.merge(graphdef, params, perturbations)
...   return jnp.mean((model(inputs) - targets) ** 2)

>>> (grads, perturbations) = grad_loss(params, perturbations, x, y)
>>> # `perturbations.xgrad.value` is the intermediate gradient
>>> assert not jnp.array_equal(perturbations.xgrad.value, jnp.zeros((1, 3)))
Parameters:

  • name (str) – A string denoting the Module attribute name for the perturbation value.

  • value (Any) – The value to take intermediate gradient.

  • variable_type (str | type[Variable[Any]]) – The Variable type for the stored perturbation. Defaulted at nnx.Perturbation.

predict(x)[source]

Makes prediction based on the input and parameters.

Parameters:

x (Array) – The input Array to make predictions on

Returns:

The calibrated probabilities

Return type:

probs

set_attributes(*filters, raise_if_not_found=True, **attributes)

Sets the attributes of nested Modules including the current Module. If the attribute is not found in the Module, it is ignored.

Example:

>>> from flax import nnx
...
>>> class Block(nnx.Module):
...   def __init__(self, din, dout, *, rngs: nnx.Rngs):
...     self.linear = nnx.Linear(din, dout, rngs=rngs)
...     self.dropout = nnx.Dropout(0.5, deterministic=False)
...     self.batch_norm = nnx.BatchNorm(10, use_running_average=False, rngs=rngs)
...
>>> block = Block(2, 5, rngs=nnx.Rngs(0))
>>> block.dropout.deterministic, block.batch_norm.use_running_average
(False, False)
>>> block.set_attributes(deterministic=True, use_running_average=True)
>>> block.dropout.deterministic, block.batch_norm.use_running_average
(True, True)

Filter’s can be used to set the attributes of specific Modules:

>>> block = Block(2, 5, rngs=nnx.Rngs(0))
>>> block.set_attributes(nnx.Dropout, deterministic=True)
>>> # Only the dropout will be modified
>>> block.dropout.deterministic, block.batch_norm.use_running_average
(True, False)
Parameters:

  • *filters (filterlib.Filter) – Filters to select the Modules to set the attributes of.

  • raise_if_not_found (bool) – If True (default), raises a ValueError if at least one attribute instance is not found in one of the selected Modules.

  • **attributes (tp.Any) – The attributes to set.

Return type:

None

sow(variable_type, name, value, reduce_fn=<function <lambda>>, init_fn=<function <lambda>>)

sow() can be used to collect intermediate values without the overhead of explicitly passing a container through each Module call. sow() stores a value in a new Module attribute, denoted by name. The value will be wrapped by a Variable of type variable_type, which can be useful to filter for in split(), state() and pop().

By default the values are stored in a tuple and each stored value is appended at the end. This way all intermediates can be tracked when the same module is called multiple times.

Example usage:

>>> from flax import nnx
>>> import jax.numpy as jnp

>>> class Model(nnx.Module):
...   def __init__(self, rngs):
...     self.linear1 = nnx.Linear(2, 3, rngs=rngs)
...     self.linear2 = nnx.Linear(3, 4, rngs=rngs)
...   def __call__(self, x, add=0):
...     x = self.linear1(x)
...     self.sow(nnx.Intermediate, 'i', x+add)
...     x = self.linear2(x)
...     return x

>>> x = jnp.ones((1, 2))
>>> model = Model(rngs=nnx.Rngs(0))
>>> assert not hasattr(model, 'i')

>>> y = model(x)
>>> assert hasattr(model, 'i')
>>> assert len(model.i.value) == 1 # tuple of length 1
>>> assert model.i.value[0].shape == (1, 3)

>>> y = model(x, add=1)
>>> assert len(model.i.value) == 2 # tuple of length 2
>>> assert (model.i.value[0] + 1 == model.i.value[1]).all()

Alternatively, a custom init/reduce function can be passed:

>>> class Model(nnx.Module):
...   def __init__(self, rngs):
...     self.linear1 = nnx.Linear(2, 3, rngs=rngs)
...     self.linear2 = nnx.Linear(3, 4, rngs=rngs)
...   def __call__(self, x):
...     x = self.linear1(x)
...     self.sow(nnx.Intermediate, 'sum', x,
...              init_fn=lambda: 0,
...              reduce_fn=lambda prev, curr: prev+curr)
...     self.sow(nnx.Intermediate, 'product', x,
...              init_fn=lambda: 1,
...              reduce_fn=lambda prev, curr: prev*curr)
...     x = self.linear2(x)
...     return x

>>> x = jnp.ones((1, 2))
>>> model = Model(rngs=nnx.Rngs(0))

>>> y = model(x)
>>> assert (model.sum.value == model.product.value).all()
>>> intermediate = model.sum.value

>>> y = model(x)
>>> assert (model.sum.value == intermediate*2).all()
>>> assert (model.product.value == intermediate**2).all()
Parameters:

  • variable_type (type[Variable[B]] | str) – The Variable type for the stored value. Typically Intermediate is used to indicate an intermediate value.

  • name (str) – A string denoting the Module attribute name, where the sowed value is stored.

  • value (A) – The value to be stored.

  • reduce_fn (Callable[[B, A], B]) – The function used to combine the existing value with the new value. The default is to append the value to a tuple.

  • init_fn (Callable[[], B]) – For the first value stored, reduce_fn will be passed the result of init_fn together with the value to be stored. The default is an empty tuple.

Return type:

bool

train(**attributes)

Sets the Module to training mode.

train uses set_attributes to recursively set attributes deterministic=False and use_running_average=False of all nested Modules that have these attributes. Its primarily used to control the runtime behavior of the Dropout and BatchNorm Modules.

Example:

>>> from flax import nnx
...
>>> class Block(nnx.Module):
...   def __init__(self, din, dout, *, rngs: nnx.Rngs):
...     self.linear = nnx.Linear(din, dout, rngs=rngs)
...     # initialize Dropout and BatchNorm in eval mode
...     self.dropout = nnx.Dropout(0.5, deterministic=True)
...     self.batch_norm = nnx.BatchNorm(10, use_running_average=True, rngs=rngs)
...
>>> block = Block(2, 5, rngs=nnx.Rngs(0))
>>> block.dropout.deterministic, block.batch_norm.use_running_average
(True, True)
>>> block.train()
>>> block.dropout.deterministic, block.batch_norm.use_running_average
(False, False)
Parameters:

**attributes – additional attributes passed to set_attributes.