mlx.nn.value_and_grad

Contents

mlx.nn.value_and_grad#

value_and_grad(model: Module, fn: Callable)#

Transform the passed function fn to a function that computes the gradients of fn wrt the model’s trainable parameters and also its value.

Parameters:

  • model (Module) – The model whose trainable parameters to compute gradients for

  • fn (Callable) – The scalar function to compute gradients for

Returns:

A callable that returns the value of fn and the gradients wrt the trainable parameters of model

Contents