Optimizer

class braintools.optim.Optimizer(*args, **kwargs)

Base class for all optimizers.

Subclasses must implement register_trainable_weights() and update() to register the parameters to optimize and to apply a single optimization step, respectively.

register_trainable_weights(param_states)

Register the trainable weights with the optimizer.

Parameters:

param_states (Dict[Hashable, State]) – The trainable weights to optimize, as a pytree whose leaves are brainstate.State objects.

Raises:

NotImplementedError – Always, in the base class. Subclasses must override this method.

update(grads)

Update the trainable weights from their gradients.

Parameters:

grads (Dict[Hashable, PyTree]) – The gradients of the loss with respect to each registered weight, with the same structure as the registered parameters.

Raises:

NotImplementedError – Always, in the base class. Subclasses must override this method.