Optimizer

Optimizer#

class seli.opt.Optimizer[source]#

Bases: Module

Base class for all gradient basedoptimizers.

Methods Summary

__call__(model, loss, grads, values)

Process the gradients of the whole model.

call_model(model, loss, grads, values)

Process the gradients of the whole model.

call_param(loss, key, grad, param)

Process the gradients of a single parameter.

minimize(loss_fn, model, *args, **kwargs)

Minimize the loss function with the given optimizer.

Methods Documentation

__call__(model: None | bool | int | float | str | type | Array | ShapeDtypeStruct | list | dict | Module | Any, loss: Float[Array, ''], grads: dict[str, Float[Array, '...']], values: dict[str, Float[Array, '...']]) → dict[str, Float[Array, '...']][source]#

Process the gradients of the whole model. The absolute loss value and parameter values are also provided to the optimizer.

Parameters:

  • model (NodeType) – The model to process.

  • loss (Float[Array, ""]) – The absolute loss value.

  • grads (dict[str, Float[Array, "..."]]) – The gradients of the model parameters.

  • values (dict[str, Float[Array, "..."]]) – The parameter values of the model.

call_model(model: None | bool | int | float | str | type | Array | ShapeDtypeStruct | list | dict | Module | Any, loss: Float[Array, ''], grads: dict[str, Float[Array, '...']], values: dict[str, Float[Array, '...']]) → dict[str, Float[Array, '...']][source]#

Process the gradients of the whole model. The absolute loss value and parameter values are also provided to the optimizer.

This function is useful for implementing custom optimizers that work on the whole model at once.

Parameters:

  • model (NodeType) – The model to process.

  • loss (Float[Array, ""]) – The absolute loss value.

  • grads (dict[str, Float[Array, "..."]]) – The gradients of the model parameters.

  • values (dict[str, Float[Array, "..."]]) – The parameter values of the model.

Returns:

grads – The processed gradients of the model parameters.

Return type:

dict[str, Float[Array, “…”]]

call_param(loss: Float[Array, ''], key: str, grad: Float[Array, '*s'], param: Float[Array, '*s']) → Float[Array, '*s'][source]#

Process the gradients of a single parameter. This function is useful for implementing custom optimizers that essentially run the same function for all parameters. This is the case for most well known optimizers.

Parameters:

  • loss (Float[Array, ""]) – The absolute loss value.

  • key (str) – The key of the parameter.

  • grad (Float[Array]) – The gradients of the parameter.

  • param (Float[Array]) – The parameter values.

Returns:

grad – The processed gradients of the parameter.

Return type:

Float[Array]

minimize(loss_fn: Loss, model: M, *args: Any, **kwargs: Any) → tuple[Self, M, Float[Array, '']][source]#

Minimize the loss function with the given optimizer.

Parameters:

  • loss_fn (Loss) – The loss function to minimize.

  • model (NodeType) – The model to minimize the loss function for.

  • args (Any) – Additional arguments to pass to the loss function.

  • kwargs (Any) – Additional keyword arguments to pass to the loss function.

Returns:

  • optimizer (Optimizer) – The optimizer.

  • model (NodeType) – The model.

  • loss (Float[Array, “”]) – The loss value.

Contents