Optimizer#

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

Base Optimizer Class.

register_trainable_weights(param_states)[source]#

Register the trainable weights.

Parameters:#

param_states: Dict[Hashable, State]

The trainable weights.

update(grads)[source]#

Update the trainable weights according to weight gradients.

Parameters:#

grads: Dict[Hashable, PyTree]

The weight gradients.