# -*- coding: utf-8 -*-
import inspect
from functools import partial, wraps
from typing import Union, Callable, Dict, Sequence, Any, Optional
import jax
import numpy as np
if jax.__version__ >= '0.4.16':
from jax.extend import linear_util
else:
from jax import linear_util
from jax import dtypes, vmap, numpy as jnp, core
from jax._src.api import (_vjp, _jvp)
from jax.api_util import argnums_partial
from jax.interpreters import xla
from jax.tree_util import (tree_flatten, tree_unflatten,
tree_map, tree_transpose,
tree_structure)
from jax.util import safe_map
from brainpy import tools, check
from brainpy._src.math.ndarray import Array, _as_jax_array_
from .tools import (dynvar_deprecation,
node_deprecation,
get_stack_cache,
cache_stack)
from .base import (BrainPyObject, ObjectTransform)
from .variables import (Variable, VariableStack)
from .tools import eval_shape
__all__ = [
'grad', # gradient of scalar function
'vector_grad', # gradient of vector/matrix/...
'functional_vector_grad',
'jacobian', 'jacrev', 'jacfwd', # gradient of jacobian
'hessian', # gradient of hessian
]
class GradientTransform(ObjectTransform):
"""Object-oriented Automatic Differentiation Transformation in BrainPy.
"""
def __init__(
self,
target: Callable,
transform: Callable,
# variables and nodes
grad_vars: Any,
dyn_vars: Dict[str, Variable],
child_objs: Dict[str, Variable],
# gradient setting
argnums: Optional[Union[int, Sequence[int]]],
return_value: bool,
has_aux: bool,
transform_setting: Optional[Dict[str, Any]] = None,
# other
name: str = None,
):
super().__init__(name=name)
# gradient variables
self._grad_vars, self._grad_tree = tree_flatten(grad_vars, is_leaf=lambda a: isinstance(a, Array))
# register variables and nodes
self.register_implicit_vars(dyn_vars, self._grad_vars)
self.register_implicit_nodes(child_objs)
# parameters
if argnums is None and len(self._grad_vars) == 0:
argnums = 0
if argnums is None:
assert len(self._grad_vars) > 0
_argnums = 0
elif isinstance(argnums, int):
_argnums = (0, argnums + 2) if len(self._grad_vars) > 0 else (argnums + 2)
else:
_argnums = check.is_sequence(argnums, elem_type=int, allow_none=False)
_argnums = tuple(a + 2 for a in _argnums)
if len(self._grad_vars) > 0:
_argnums = (0,) + _argnums
self._nonvar_argnums = argnums
self._argnums = _argnums
self._return_value = return_value
self._has_aux = has_aux
# target
self.target = target
# transform
self._eval_dyn_vars = False
self._grad_transform = transform
self._dyn_vars = VariableStack()
self._transform = None
self._grad_setting = dict() if transform_setting is None else transform_setting
if self._has_aux:
self._transform = self._grad_transform(
self._f_grad_with_aux_to_transform,
argnums=self._argnums,
has_aux=True,
**self._grad_setting
)
else:
self._transform = self._grad_transform(
self._f_grad_without_aux_to_transform,
argnums=self._argnums,
has_aux=True,
**self._grad_setting
)
def _f_grad_with_aux_to_transform(self,
grad_values: tuple,
dyn_values: dict,
*args,
**kwargs):
for k in dyn_values.keys():
self._dyn_vars[k]._value = dyn_values[k]
for v, d in zip(self._grad_vars, grad_values):
v._value = d
# Users should return the auxiliary data like::
# >>> # 1. example of return one data
# >>> return scalar_loss, data
# >>> # 2. example of return multiple data
# >>> return scalar_loss, (data1, data2, ...)
outputs = self.target(*args, **kwargs)
# outputs: [0] is the value for gradient,
# [1] is other values for return
output0 = tree_map(lambda a: (a.value if isinstance(a, Array) else a), outputs[0])
return output0, (outputs, [v.value for v in self._grad_vars], self._dyn_vars.dict_data())
def _f_grad_without_aux_to_transform(self,
grad_values: tuple,
dyn_values: dict,
*args,
**kwargs):
for k in dyn_values.keys():
self._dyn_vars[k]._value = dyn_values[k]
for v, d in zip(self._grad_vars, grad_values):
v._value = d
# Users should return the scalar value like this::
# >>> return scalar_loss
output = self.target(*args, **kwargs)
output0 = tree_map(lambda a: (a.value if isinstance(a, Array) else a), output)
return output0, (output, [v.value for v in self._grad_vars], self._dyn_vars.dict_data())
def __repr__(self):
name = self.__class__.__name__
f = tools.repr_object(self.target)
f = tools.repr_context(f, " " * (len(name) + 6))
format_ref = (f'{name}({self.name}, target={f}, \n' +
f'{" " * len(name)} num_of_grad_vars={len(self._grad_vars)}, \n'
f'{" " * len(name)} num_of_dyn_vars={len(self._dyn_vars)})')
return format_ref
def _return(self, rets):
grads, (outputs, new_grad_vs, new_dyn_vs) = rets
for v, d in zip(self._grad_vars, new_grad_vs):
v._value = d
for k in new_dyn_vs.keys():
self._dyn_vars[k]._value = new_dyn_vs[k]
# check returned grads
if len(self._grad_vars) > 0:
if self._nonvar_argnums is None:
grads = self._grad_tree.unflatten(grads)
else:
var_grads = self._grad_tree.unflatten(grads[0])
arg_grads = grads[1] if isinstance(self._nonvar_argnums, int) else grads[1:]
grads = (var_grads, arg_grads)
# check returned value
if self._return_value:
# check aux
if self._has_aux:
return grads, outputs[0], outputs[1]
else:
return grads, outputs
else:
# check aux
if self._has_aux:
return grads, outputs[1]
else:
return grads
def __call__(self, *args, **kwargs):
if jax.config.jax_disable_jit: # disable JIT
rets = self._transform(
[v.value for v in self._grad_vars], # variables for gradients
self._dyn_vars.dict_data(), # dynamical variables
*args,
**kwargs
)
return self._return(rets)
elif not self._eval_dyn_vars: # evaluate dynamical variables
stack = get_stack_cache(self.target)
if stack is None:
with VariableStack() as stack:
rets = eval_shape(self._transform,
[v.value for v in self._grad_vars], # variables for gradients
{}, # dynamical variables
*args,
**kwargs)
cache_stack(self.target, stack)
self._dyn_vars = stack
self._dyn_vars.remove_by_id(*[id(v) for v in self._grad_vars])
self._eval_dyn_vars = True
# if not the outermost transformation
if not stack.is_first_stack():
return self._return(rets)
rets = self._transform(
[v.value for v in self._grad_vars], # variables for gradients
self._dyn_vars.dict_data(), # dynamical variables
*args,
**kwargs
)
return self._return(rets)
def _make_grad(
func: Callable,
grad_vars: Optional[Union[Variable, Sequence[Variable], Dict[str, Variable]]] = None,
argnums: Optional[Union[int, Sequence[int]]] = None,
holomorphic: Optional[bool] = False,
allow_int: Optional[bool] = False,
reduce_axes: Optional[Sequence[str]] = (),
has_aux: Optional[bool] = None,
return_value: Optional[bool] = False,
# deprecated
dyn_vars: Optional[Union[Variable, Sequence[Variable], Dict[str, Variable]]] = None,
child_objs: Optional[Union[BrainPyObject, Sequence[BrainPyObject], Dict[str, BrainPyObject]]] = None,
):
child_objs = check.is_all_objs(child_objs, out_as='dict')
dyn_vars = check.is_all_vars(dyn_vars, out_as='dict')
return GradientTransform(target=func,
transform=jax.grad,
grad_vars=grad_vars,
dyn_vars=dyn_vars,
child_objs=child_objs,
argnums=argnums,
return_value=return_value,
has_aux=False if has_aux is None else has_aux,
transform_setting=dict(holomorphic=holomorphic,
allow_int=allow_int,
reduce_axes=reduce_axes))
[docs]
def grad(
func: Optional[Callable] = None,
grad_vars: Optional[Union[Variable, Sequence[Variable], Dict[str, Variable]]] = None,
argnums: Optional[Union[int, Sequence[int]]] = None,
holomorphic: Optional[bool] = False,
allow_int: Optional[bool] = False,
reduce_axes: Optional[Sequence[str]] = (),
has_aux: Optional[bool] = None,
return_value: Optional[bool] = False,
# deprecated
dyn_vars: Optional[Union[Variable, Sequence[Variable], Dict[str, Variable]]] = None,
child_objs: Optional[Union[BrainPyObject, Sequence[BrainPyObject], Dict[str, BrainPyObject]]] = None,
) -> Union[Callable, GradientTransform]:
"""Automatic gradient computation for functions or class objects.
This gradient function only support scalar return. It creates a function
which evaluates the gradient of ``func``.
It's worthy to note that the returns are different for different argument settings (where ``arg_grads`` refers
to the gradients of "argnums", and ``var_grads`` refers to the gradients of "grad_vars").
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)``.
2. 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)``.
3. 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)``.
Let's see some examples below.
Before start, let's figure out what should be provided as ``grad_vars``?
And, what should be labeled in ``argnums``?
Take the following codes as example:
>>> import brainpy as bp
>>> import brainpy.math as bm
>>>
>>> class Example(bp.BrainPyObject):
>>> def __init__(self):
>>> super(Example, self).__init__()
>>> self.x = bm.TrainVar(bm.zeros(1))
>>> self.y = bm.random.rand(10)
>>> def __call__(self, z, v):
>>> t1 = self.x * self.y.sum()
>>> t2 = bm.tanh(z * v + t1)
>>> return t2.mean()
>>>
>>> # This code is equivalent to the following function:
>>>
>>> x = bm.TrainVar(bm.zeros(1))
>>> y = bm.random.rand(10)
>>> def f(z, v):
>>> t1 = x * y.sum()
>>> t2 = bm.tanh(z * v + t1)
>>> return t2.mean()
Generally speaking, all gradient variables which not provided in arguments should be
labeled as ``grad_vars``, while all gradient variables provided in the function arguments
should be declared in ``argnums``.
In above codes, we try to take gradients of ``self.x`` and arguments ``z`` and ``v``, we should
call ``brainpy.math.grad`` as:
>>> f = Example()
>>> f_grad = bm.grad(f, grad_vars=f.x, argnums=(0, 1))
Examples
--------
Grad for a pure function:
>>> import brainpy as bp
>>> grad_tanh = grad(bp.math.tanh)
>>> print(grad_tanh(0.2))
0.961043
Parameters
----------
func : callable, function, BrainPyObject
Function to be differentiated. Its arguments at positions specified by
``argnums`` should be arrays, scalars, or standard Python containers.
Argument arrays in the positions specified by ``argnums`` must be of
inexact (i.e., floating-point or complex) type. It should return a scalar
(which includes arrays with shape ``()`` but not arrays with shape ``(1,)`` etc.)
grad_vars : optional, ArrayType, sequence of ArrayType, dict
The variables in ``func`` to take their gradients.
argnums : optional, integer or sequence of integers
Specifies which positional argument(s) to differentiate with respect to (default 0).
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.
holomorphic: optional, bool
Indicates whether ``fun`` is promised to be
holomorphic. If True, inputs and outputs must be complex. Default False.
allow_int: optional, bool
Whether to allow differentiating with
respect to integer valued inputs. The gradient of an integer input will
have a trivial vector-space dtype (float0). Default False.
reduce_axes: optional, tuple of int
tuple of axis names. If an axis is listed here, and
``fun`` implicitly broadcasts a value over that axis, the backward pass
will perform a ``psum`` of the corresponding gradient. Otherwise, the
gradient will be per-example over named axes. For example, if ``'batch'``
is a named batch axis, ``grad(f, reduce_axes=('batch',))`` will create a
function that computes the total gradient while ``grad(f)`` will create
one that computes the per-example gradient.
dyn_vars : optional, ArrayType, sequence of ArrayType, dict
The dynamically changed variables used in ``func``.
.. deprecated:: 2.4.0
No longer need to provide ``dyn_vars``. This function is capable of automatically
collecting the dynamical variables used in the target ``func``.
child_objs: optional, BrainPyObject, sequnce, dict
.. versionadded:: 2.3.1
.. deprecated:: 2.4.0
No longer need to provide ``child_objs``. This function is capable of automatically
collecting the children objects used in the target ``func``.
Returns
-------
func : GradientTransform
A function with the same arguments as ``fun``, that evaluates the gradient
of ``fun``. If ``argnums`` is an integer then the gradient has the same
shape and type as the positional argument indicated by that integer. If
argnums is a tuple of integers, the gradient is a tuple of values with the
same shapes and types as the corresponding arguments. If ``has_aux`` is True
then a pair of (gradient, auxiliary_data) is returned.
"""
dynvar_deprecation(dyn_vars)
node_deprecation(child_objs)
if func is None:
return lambda f: _make_grad(f,
grad_vars=grad_vars,
dyn_vars=dyn_vars,
child_objs=child_objs,
argnums=argnums,
holomorphic=holomorphic,
allow_int=allow_int,
reduce_axes=reduce_axes,
has_aux=has_aux,
return_value=return_value)
else:
return _make_grad(func=func,
grad_vars=grad_vars,
dyn_vars=dyn_vars,
child_objs=child_objs,
argnums=argnums,
holomorphic=holomorphic,
allow_int=allow_int,
reduce_axes=reduce_axes,
has_aux=has_aux,
return_value=return_value)
def _unravel_array_into_pytree(pytree, axis, arr, is_leaf=None):
leaves, treedef = tree_flatten(pytree, is_leaf=is_leaf)
axis = axis % arr.ndim
shapes = [arr.shape[:axis] + np.shape(l) + arr.shape[axis + 1:] for l in leaves]
parts = jnp.split(_as_jax_array_(arr), np.cumsum(safe_map(np.size, leaves[:-1])), axis)
reshaped_parts = [x.reshape(shape) for x, shape in zip(parts, shapes)]
return tree_unflatten(treedef, reshaped_parts, )
def _std_basis(pytree):
leaves, _ = tree_flatten(pytree)
ndim = sum(safe_map(np.size, leaves))
dtype = dtypes.result_type(*leaves)
flat_basis = jax.numpy.eye(ndim, dtype=dtype)
return _unravel_array_into_pytree(pytree, 1, flat_basis)
def _isleaf(x):
return isinstance(x, Array)
def _jacrev(fun, argnums=0, holomorphic=False, allow_int=False, has_aux=False, return_value=False):
_check_callable(fun)
@wraps(fun)
def jacfun(*args, **kwargs):
f = linear_util.wrap_init(fun, kwargs)
f_partial, dyn_args = argnums_partial(f, argnums, args, require_static_args_hashable=False)
tree_map(partial(_check_input_dtype_jacrev, holomorphic, allow_int), dyn_args)
if has_aux:
y, pullback, aux = _vjp(f_partial, *dyn_args, has_aux=True)
else:
y, pullback = _vjp(f_partial, *dyn_args, has_aux=False)
tree_map(partial(_check_output_dtype_jacrev, holomorphic), y)
jac = vmap(pullback)(_std_basis(y))
jac = jac[0] if isinstance(argnums, int) else jac
example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args
jac_tree = tree_map(partial(_unravel_array_into_pytree, y, 0, is_leaf=_isleaf), jac, is_leaf=_isleaf)
jac = tree_transpose(tree_structure(example_args), tree_flatten(y, is_leaf=_isleaf)[1], jac_tree)
if return_value:
return (jac, y, aux) if has_aux else (jac, y)
else:
return (jac, aux) if has_aux else jac
return jacfun
[docs]
def jacrev(
func: Callable,
grad_vars: Optional[Union[Variable, Sequence[Variable], Dict[str, Variable]]] = None,
argnums: Optional[Union[int, Sequence[int]]] = None,
has_aux: Optional[bool] = None,
return_value: bool = False,
holomorphic: bool = False,
allow_int: bool = False,
# deprecated
dyn_vars: Optional[Union[Variable, Sequence[Variable], Dict[str, Variable]]] = None,
child_objs: Optional[Union[BrainPyObject, Sequence[BrainPyObject], Dict[str, BrainPyObject]]] = None,
) -> ObjectTransform:
"""Extending automatic Jacobian (reverse-mode) of ``func`` to classes.
This function extends the JAX official ``jacrev`` to make automatic jacobian
computation on functions and class functions. Moreover, it supports returning
value ("return_value") and returning auxiliary data ("has_aux").
Same as `brainpy.math.grad <./brainpy.math.autograd.grad.html>`_, the returns are
different for different argument settings in ``brainpy.math.jacrev``.
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)``.
2. 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)``.
3. 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: Function whose Jacobian 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``).
holomorphic: Optional, bool.
Indicates whether ``fun`` is promised to be
holomorphic. Default False.
allow_int: Optional, bool.
Whether to allow differentiating with
respect to integer valued inputs. The gradient of an integer input will
have a trivial vector-space dtype (float0). Default False.
dyn_vars : optional, ArrayType, sequence of ArrayType, dict
The dynamically changed variables used in ``func``.
.. deprecated:: 2.4.0
No longer need to provide ``dyn_vars``. This function is capable of automatically
collecting the dynamical variables used in the target ``func``.
child_objs: optional, BrainPyObject, sequnce, dict
.. versionadded:: 2.3.1
.. deprecated:: 2.4.0
No longer need to provide ``child_objs``. This function is capable of automatically
collecting the children objects used in the target ``func``.
Returns
-------
fun: GradientTransform
The transformed object.
"""
child_objs = check.is_all_objs(child_objs, out_as='dict')
dyn_vars = check.is_all_vars(dyn_vars, out_as='dict')
return GradientTransform(target=func,
transform=_jacrev,
grad_vars=grad_vars,
dyn_vars=dyn_vars,
child_objs=child_objs,
argnums=argnums,
return_value=return_value,
has_aux=False if has_aux is None else has_aux,
transform_setting=dict(holomorphic=holomorphic,
allow_int=allow_int))
jacobian = jacrev
def _jacfwd(fun, argnums=0, holomorphic=False, has_aux=False, return_value=False):
_check_callable(fun)
@wraps(fun)
def jacfun(*args, **kwargs):
f = linear_util.wrap_init(fun, kwargs)
f_partial, dyn_args = argnums_partial(f, argnums, args, require_static_args_hashable=False)
tree_map(partial(_check_input_dtype_jacfwd, holomorphic), dyn_args)
if has_aux:
pushfwd = partial(_jvp, f_partial, dyn_args, has_aux=True)
y, jac, aux = vmap(pushfwd, out_axes=(None, -1, None))(_std_basis(dyn_args))
else:
pushfwd = partial(_jvp, f_partial, dyn_args)
y, jac = vmap(pushfwd, out_axes=(None, -1))(_std_basis(dyn_args))
tree_map(partial(_check_output_dtype_jacfwd, holomorphic), y)
example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args
jac = tree_map(partial(_unravel_array_into_pytree, example_args, -1, is_leaf=_isleaf), jac, is_leaf=_isleaf)
if return_value:
return (jac, y, aux) if has_aux else (jac, y)
else:
return (jac, aux) if has_aux else jac
return jacfun
[docs]
def jacfwd(
func: Callable,
grad_vars: Optional[Union[Variable, Sequence[Variable], Dict[str, Variable]]] = None,
argnums: Optional[Union[int, Sequence[int]]] = None,
has_aux: Optional[bool] = None,
return_value: bool = False,
holomorphic: bool = False,
# deprecated
dyn_vars: Optional[Union[Variable, Sequence[Variable], Dict[str, Variable]]] = None,
child_objs: Optional[Union[BrainPyObject, Sequence[BrainPyObject], Dict[str, BrainPyObject]]] = None,
) -> ObjectTransform:
"""Extending automatic Jacobian (forward-mode) of ``func`` to classes.
This function extends the JAX official ``jacfwd`` to make automatic jacobian
computation on functions and class functions. Moreover, it supports returning
value ("return_value") and returning auxiliary data ("has_aux").
Same as `brainpy.math.grad <./brainpy.math.autograd.grad.html>`_, the returns are
different for different argument settings in ``brainpy.math.jacfwd``.
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)``.
2. 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)``.
3. 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: Function whose Jacobian 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``).
holomorphic: Optional, bool. Indicates whether ``fun`` is promised to be
holomorphic. Default False.
dyn_vars : optional, ArrayType, sequence of ArrayType, dict
The dynamically changed variables used in ``func``.
.. deprecated:: 2.4.0
No longer need to provide ``dyn_vars``. This function is capable of automatically
collecting the dynamical variables used in the target ``func``.
child_objs: optional, BrainPyObject, sequnce, dict
.. versionadded:: 2.3.1
.. deprecated:: 2.4.0
No longer need to provide ``child_objs``. This function is capable of automatically
collecting the children objects used in the target ``func``.
Returns
-------
obj: GradientTransform
The transformed object.
"""
child_objs = check.is_all_objs(child_objs, out_as='dict')
dyn_vars = check.is_all_vars(dyn_vars, out_as='dict')
return GradientTransform(target=func,
transform=_jacfwd,
grad_vars=grad_vars,
dyn_vars=dyn_vars,
child_objs=child_objs,
argnums=argnums,
return_value=return_value,
has_aux=False if has_aux is None else has_aux,
transform_setting=dict(holomorphic=holomorphic))
def _functional_hessian(
fun: Callable,
argnums: Optional[Union[int, Sequence[int]]] = None,
has_aux: bool = False,
holomorphic: bool = False,
):
return _jacfwd(
_jacrev(fun, argnums, has_aux=has_aux, holomorphic=holomorphic),
argnums, has_aux=has_aux, holomorphic=holomorphic
)
class GradientTransformPreserveTree(ObjectTransform):
"""
Object-oriented Automatic Differentiation Transformation in BrainPy.
"""
def __init__(
self,
target: Callable,
transform: Callable,
# variables and nodes
grad_vars: Dict[str, Variable],
# gradient setting
argnums: Optional[Union[int, Sequence[int]]],
return_value: bool,
has_aux: bool,
transform_setting: Optional[Dict[str, Any]] = None,
# other
name: str = None,
):
super().__init__(name=name)
# gradient variables
if grad_vars is None:
grad_vars = dict()
assert isinstance(grad_vars, dict), 'grad_vars should be a dict'
new_grad_vars = {}
for k, v in grad_vars.items():
assert isinstance(v, Variable)
new_grad_vars[k] = v
self._grad_vars = new_grad_vars
# parameters
if argnums is None and len(self._grad_vars) == 0:
argnums = 0
if argnums is None:
assert len(self._grad_vars) > 0
_argnums = 0
elif isinstance(argnums, int):
_argnums = (0, argnums + 2) if len(self._grad_vars) > 0 else (argnums + 2)
else:
_argnums = check.is_sequence(argnums, elem_type=int, allow_none=False)
_argnums = tuple(a + 2 for a in _argnums)
if len(self._grad_vars) > 0:
_argnums = (0,) + _argnums
self._nonvar_argnums = argnums
self._argnums = _argnums
self._return_value = return_value
self._has_aux = has_aux
# target
self.target = target
# transform
self._eval_dyn_vars = False
self._grad_transform = transform
self._dyn_vars = VariableStack()
self._transform = None
self._grad_setting = dict() if transform_setting is None else transform_setting
if self._has_aux:
self._transform = self._grad_transform(
self._f_grad_with_aux_to_transform,
argnums=self._argnums,
has_aux=True,
**self._grad_setting
)
else:
self._transform = self._grad_transform(
self._f_grad_without_aux_to_transform,
argnums=self._argnums,
has_aux=True,
**self._grad_setting
)
def _f_grad_with_aux_to_transform(self,
grad_values: dict,
dyn_values: dict,
*args,
**kwargs):
for k in dyn_values.keys():
self._dyn_vars[k]._value = dyn_values[k]
for k, v in grad_values.items():
self._grad_vars[k]._value = v
# Users should return the auxiliary data like::
# >>> # 1. example of return one data
# >>> return scalar_loss, data
# >>> # 2. example of return multiple data
# >>> return scalar_loss, (data1, data2, ...)
outputs = self.target(*args, **kwargs)
# outputs: [0] is the value for gradient,
# [1] is other values for return
output0 = tree_map(lambda a: (a.value if isinstance(a, Array) else a), outputs[0])
return output0, (outputs, {k: v for k, v in self._grad_vars.items()}, self._dyn_vars.dict_data())
def _f_grad_without_aux_to_transform(self,
grad_values: dict,
dyn_values: dict,
*args,
**kwargs):
for k in dyn_values.keys():
self._dyn_vars[k].value = dyn_values[k]
for k, v in grad_values.items():
self._grad_vars[k].value = v
# Users should return the scalar value like this::
# >>> return scalar_loss
output = self.target(*args, **kwargs)
output0 = tree_map(lambda a: (a.value if isinstance(a, Array) else a), output)
return output0, (output, {k: v.value for k, v in self._grad_vars.items()}, self._dyn_vars.dict_data())
def __repr__(self):
name = self.__class__.__name__
f = tools.repr_object(self.target)
f = tools.repr_context(f, " " * (len(name) + 6))
format_ref = (f'{name}({self.name}, target={f}, \n' +
f'{" " * len(name)} num_of_grad_vars={len(self._grad_vars)}, \n'
f'{" " * len(name)} num_of_dyn_vars={len(self._dyn_vars)})')
return format_ref
def _return(self, rets):
grads, (outputs, new_grad_vs, new_dyn_vs) = rets
for k, v in new_grad_vs.items():
self._grad_vars[k].value = v
for k in new_dyn_vs.keys():
self._dyn_vars[k].value = new_dyn_vs[k]
# check returned grads
if len(self._grad_vars) > 0:
if self._nonvar_argnums is None:
pass
else:
arg_grads = grads[1] if isinstance(self._nonvar_argnums, int) else grads[1:]
grads = (grads[0], arg_grads)
# check returned value
if self._return_value:
# check aux
if self._has_aux:
return grads, outputs[0], outputs[1]
else:
return grads, outputs
else:
# check aux
if self._has_aux:
return grads, outputs[1]
else:
return grads
def __call__(self, *args, **kwargs):
if jax.config.jax_disable_jit: # disable JIT
rets = self._transform(
{k: v.value for k, v in self._grad_vars.items()}, # variables for gradients
self._dyn_vars.dict_data(), # dynamical variables
*args,
**kwargs
)
return self._return(rets)
elif not self._eval_dyn_vars: # evaluate dynamical variables
stack = get_stack_cache(self.target)
if stack is None:
with VariableStack() as stack:
rets = eval_shape(
self._transform,
{k: v.value for k, v in self._grad_vars.items()}, # variables for gradients
{}, # dynamical variables
*args,
**kwargs
)
cache_stack(self.target, stack)
self._dyn_vars = stack
self._dyn_vars.remove_by_id(*[id(v) for v in self._grad_vars.values()])
self._eval_dyn_vars = True
# if not the outermost transformation
if not stack.is_first_stack():
return self._return(rets)
rets = self._transform(
{k: v.value for k, v in self._grad_vars.items()}, # variables for gradients
self._dyn_vars.dict_data(), # dynamical variables
*args,
**kwargs
)
return self._return(rets)
[docs]
def hessian(
func: Callable,
grad_vars: Optional[Union[Variable, Sequence[Variable], Dict[str, Variable]]] = None,
argnums: Optional[Union[int, Sequence[int]]] = None,
has_aux: Optional[bool] = None,
holomorphic=False,
) -> ObjectTransform:
"""Hessian of ``func`` as a dense array.
Parameters
----------
func : callable, function
Function whose Hessian is to be computed. Its arguments at positions
specified by ``argnums`` should be arrays, scalars, or standard Python
containers thereof. It should return arrays, scalars, or standard Python
containers thereof.
grad_vars : optional, ArrayCollector, sequence of ArrayType
The variables required to compute their gradients.
argnums: Optional, integer or sequence of integers
Specifies which positional argument(s) to differentiate with respect to (default ``0``).
holomorphic : bool
Indicates whether ``fun`` is promised to be holomorphic. Default False.
has_aux : 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.
Returns
-------
obj: ObjectTransform
The transformed object.
"""
return GradientTransformPreserveTree(target=func,
transform=jax.hessian,
grad_vars=grad_vars,
argnums=argnums,
has_aux=False if has_aux is None else has_aux,
transform_setting=dict(holomorphic=holomorphic),
return_value=False)
def functional_vector_grad(func, argnums=0, return_value=False, has_aux=False):
_check_callable(func)
@wraps(func)
def grad_fun(*args, **kwargs):
f = linear_util.wrap_init(func, kwargs)
f_partial, dyn_args = argnums_partial(f, argnums, args, require_static_args_hashable=False)
if has_aux:
y, vjp_fn, aux = _vjp(f_partial, *dyn_args, has_aux=True)
else:
y, vjp_fn = _vjp(f_partial, *dyn_args, has_aux=False)
leaves, tree = tree_flatten(y)
tangents = tree_unflatten(tree, [jnp.ones(l.shape, dtype=l.dtype) for l in leaves])
grads = vjp_fn(tangents)
if isinstance(argnums, int):
grads = grads[0]
if has_aux:
return (grads, y, aux) if return_value else (grads, aux)
else:
return (grads, y) if return_value else grads
return grad_fun
[docs]
def vector_grad(
func: Optional[Callable] = None,
grad_vars: Optional[Union[Variable, Sequence[Variable], Dict[str, Variable]]] = None,
argnums: Optional[Union[int, Sequence[int]]] = None,
return_value: bool = False,
has_aux: Optional[bool] = None,
# deprecated
dyn_vars: Optional[Union[Variable, Sequence[Variable], Dict[str, Variable]]] = None,
child_objs: Optional[Union[BrainPyObject, Sequence[BrainPyObject], Dict[str, BrainPyObject]]] = None,
) -> Union[Callable, ObjectTransform]:
"""Take vector-valued gradients for function ``func``.
Same as `brainpy.math.grad <./brainpy.math.autograd.grad.html>`_,
`brainpy.math.jacrev <./brainpy.math.autograd.jacrev.html>`_ and
`brainpy.math.jacfwd <./brainpy.math.autograd.jacfwd.html>`_,
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)``.
2. 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)``.
3. 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``).
dyn_vars : optional, ArrayType, sequence of ArrayType, dict
The dynamically changed variables used in ``func``.
.. deprecated:: 2.4.0
No longer need to provide ``dyn_vars``. This function is capable of automatically
collecting the dynamical variables used in the target ``func``.
child_objs: optional, BrainPyObject, sequnce, dict
.. versionadded:: 2.3.1
.. deprecated:: 2.4.0
No longer need to provide ``child_objs``. This function is capable of automatically
collecting the children objects used in the target ``func``.
Returns
-------
func : GradientTransform
The vector gradient function.
"""
child_objs = check.is_all_objs(child_objs, out_as='dict')
dyn_vars = check.is_all_vars(dyn_vars, out_as='dict')
if func is None:
return lambda f: GradientTransform(target=f,
transform=functional_vector_grad,
grad_vars=grad_vars,
dyn_vars=dyn_vars,
child_objs=child_objs,
argnums=argnums,
return_value=return_value,
has_aux=False if has_aux is None else has_aux)
else:
return GradientTransform(target=func,
transform=functional_vector_grad,
grad_vars=grad_vars,
dyn_vars=dyn_vars,
child_objs=child_objs,
argnums=argnums,
return_value=return_value,
has_aux=False if has_aux is None else has_aux)
def _check_callable(fun):
# In Python 3.10+, the only thing stopping us from supporting staticmethods
# is that we can't take weak references to them, which the C++ JIT requires.
if isinstance(fun, staticmethod):
raise TypeError(f"staticmethod arguments are not supported, got {fun}")
if not callable(fun):
raise TypeError(f"Expected a callable value, got {fun}")
if _isgeneratorfunction(fun):
raise TypeError(f"Expected a function, got a generator function: {fun}")
def _isgeneratorfunction(fun):
# re-implemented here because of https://bugs.python.org/issue33261
while inspect.ismethod(fun):
fun = fun.__func__
while isinstance(fun, partial):
fun = fun.func
return inspect.isfunction(fun) and bool(fun.__code__.co_flags & inspect.CO_GENERATOR)
def _check_arg(arg):
if not (isinstance(arg, core.Tracer) or _valid_jaxtype(arg)):
raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid JAX type.")
def _valid_jaxtype(arg):
try:
xla.abstractify(arg) # faster than core.get_aval
except TypeError:
return core.valid_jaxtype(arg)
else:
return True
def _check_output_dtype_revderiv(name, holomorphic, x):
aval = core.get_aval(x)
# if jnp.issubdtype(aval.dtype, dtypes.extended):
# raise TypeError(f"{name} with output element type {aval.dtype.name}")
if holomorphic:
if not dtypes.issubdtype(aval.dtype, np.complexfloating):
raise TypeError(f"{name} with holomorphic=True requires outputs with complex dtype, "
f"but got {aval.dtype.name}.")
elif dtypes.issubdtype(aval.dtype, np.complexfloating):
raise TypeError(f"{name} requires real-valued outputs (output dtype that is "
f"a sub-dtype of np.floating), but got {aval.dtype.name}. "
"For holomorphic differentiation, pass holomorphic=True. "
"For differentiation of non-holomorphic functions involving complex "
"outputs, use jax.vjp directly.")
elif not dtypes.issubdtype(aval.dtype, np.floating):
raise TypeError(f"{name} requires real-valued outputs (output dtype that is "
f"a sub-dtype of np.floating), but got {aval.dtype.name}. "
"For differentiation of functions with integer outputs, use "
"jax.vjp directly.")
def _check_input_dtype_revderiv(name, holomorphic, allow_int, x):
_check_arg(x)
aval = core.get_aval(x)
# if jnp.issubdtype(aval.dtype, dtypes.extended):
# raise TypeError(f"{name} with input element type {aval.dtype.name}")
if holomorphic:
if not dtypes.issubdtype(aval.dtype, np.complexfloating):
raise TypeError(f"{name} with holomorphic=True requires inputs with complex dtype, "
f"but got {aval.dtype.name}.")
if (dtypes.issubdtype(aval.dtype, np.integer) or
dtypes.issubdtype(aval.dtype, np.bool_)):
if not allow_int:
raise TypeError(f"{name} requires real- or complex-valued inputs (input dtype "
f"that is a sub-dtype of np.inexact), but got {aval.dtype.name}. "
"If you want to use Boolean- or integer-valued inputs, use vjp "
"or set allow_int to True.")
elif not dtypes.issubdtype(aval.dtype, np.inexact):
raise TypeError(f"{name} requires numerical-valued inputs (input dtype that is a "
f"sub-dtype of np.bool_ or np.number), but got {aval.dtype.name}.")
_check_output_dtype_jacrev = partial(_check_output_dtype_revderiv, "jacrev")
_check_input_dtype_jacrev = partial(_check_input_dtype_revderiv, "jacrev")
def _check_output_dtype_jacfwd(holomorphic, x):
aval = core.get_aval(x)
if holomorphic:
if not dtypes.issubdtype(aval.dtype, np.complexfloating):
raise TypeError("jacfwd with holomorphic=True requires outputs with complex dtype, "
f"but got {aval.dtype.name}.")
def _check_input_dtype_jacfwd(holomorphic: bool, x: Any) -> None:
_check_arg(x)
aval = core.get_aval(x)
# if jnp.issubdtype(aval.dtype, dtypes.extended):
# raise TypeError(f"jacfwd with input element type {aval.dtype.name}")
if holomorphic:
if not dtypes.issubdtype(aval.dtype, np.complexfloating):
raise TypeError("jacfwd with holomorphic=True requires inputs with complex "
f"dtype, but got {aval.dtype.name}.")
elif not dtypes.issubdtype(aval.dtype, np.floating):
raise TypeError("jacfwd requires real-valued inputs (input dtype that is "
f"a sub-dtype of np.floating), but got {aval.dtype.name}. "
"For holomorphic differentiation, pass holomorphic=True. "
"For differentiation of non-holomorphic functions involving "
"complex inputs or integer inputs, use jax.jvp directly.")
