braincore.transform.grad

braincore.transform.grad(fun=None, grad_vars=None, argnums=None, holomorphic=False, allow_int=False, reduce_axes=(), has_aux=None, return_value=False)

Compute the gradient of a scalar-valued function with respect to its arguments.

Parameters:

• reduce_axes (Optional[Sequence[str]])

• allow_int (Optional[bool])

• holomorphic (Optional[bool])

• grad_vars (Union[State, Sequence[State], Dict[str, State], None])

• fun (Optional[Callable]) – the scalar-valued function to be differentiated.

• argnums (Union[int, Sequence[int], None]) – (int or tuple of ints) optional. Specifies which positional argument(s) to differentiate with respect to.

• has_aux (Optional[bool]) – (bool) optional. 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 (Optional[bool]) – (bool) optional. Indicates whether to return the value of the function along with the gradient. Default False.

Return type:

Union[GradientTransform, Callable[[Callable], GradientTransform]]

Returns:

A function which computes the gradient of fun. The function takes the same arguments as fun, but returns the gradient instead. If has_aux is True, the function returns a pair where the first element is the gradient and the second element is the auxiliary data. If return_value is True, the function returns a pair where the first element is the gradient and the second element is the value of the function.