Ctrl+K

braincore.transform.vector_grad

Contents

braincore.transform.vector_grad#

braincore.transform.vector_grad(func=None, grad_vars=None, argnums=None, return_value=False, has_aux=None, state_auto_finding=True)#

Take vector-valued gradients for function func.

Same as brainpy.math.grad, brainpy.math.jacrev and brainpy.math.jacfwd, the returns in this function are different for different argument settings.

  1. When “grad_vars” is None

  • “has_aux=False” + “return_value=False” => arg_grads.

  • “has_aux=True” + “return_value=False” => (arg_grads, aux_data).

  • “has_aux=False” + “return_value=True” => (arg_grads, loss_value).

  • “has_aux=True” + “return_value=True” => (arg_grads, loss_value, aux_data).

  1. When “grad_vars” is not None and “argnums” is None

  • “has_aux=False” + “return_value=False” => var_grads.

  • “has_aux=True” + “return_value=False” => (var_grads, aux_data).

  • “has_aux=False” + “return_value=True” => (var_grads, loss_value).

  • “has_aux=True” + “return_value=True” => (var_grads, loss_value, aux_data).

  1. When “grad_vars” is not None and “argnums” is not None

  • “has_aux=False” + “return_value=False” => (var_grads, arg_grads).

  • “has_aux=True” + “return_value=False” => ((var_grads, arg_grads), aux_data).

  • “has_aux=False” + “return_value=True” => ((var_grads, arg_grads), loss_value).

  • “has_aux=True” + “return_value=True” => ((var_grads, arg_grads), loss_value, aux_data).

Parameters:

  • func (Callable) – Function whose gradient is to be computed.

  • grad_vars (optional, ArrayType, sequence of ArrayType, dict) – The variables in func to take their gradients.

  • has_aux (optional, bool) – Indicates whether fun returns a pair where the first element is considered the output of the mathematical function to be differentiated and the second element is auxiliary data. Default False.

  • return_value (bool) – Whether return the loss value.

  • argnums (Optional, integer or sequence of integers. Specifies which) – positional argument(s) to differentiate with respect to (default 0).

  • state_auto_finding (bool) – Whether to automatically find the states in the function.

Returns:

func – The vector gradient function.

Return type:

GradientTransform

Contents