Gradient

Computation of higher order derivatives with torch gradient function

ndmap.gradient.derivative(series: dict[tuple[int, ...], torch.Tensor], index: tuple[int, ...]) → dict[tuple[int, ...], torch.Tensor]

Compute series derivative

Parameters:

• series (dict[tuple[int, ...], Tensor]) – series

• index – derivative index

Return type:

dict[tuple[int, …], Tensor]

Examples

>>> from pprint import pprint
>>> import torch
>>> def fn(x):
...    return 1.0 + x + x**2 + x**3 + x**4 + x**5
>>> x = torch.tensor([0.0])
>>> s = series((5, ), fn, x, retain=False)
>>> pprint(derivative(s, (1, )), sort_dicts=False)
{(0,): tensor([1.]),
 (1,): tensor([2.]),
 (2,): tensor([3.]),
 (3,): tensor([4.]),
 (4,): tensor([5.])}

ndmap.gradient.evaluate(series: dict[tuple[int, ...], torch.Tensor], delta: list[torch.Tensor]) → torch.Tensor

Evaluate series

Parameters:

• series (Series) – input series representation

• delta (list[Tensor]) – delta deviation

Return type:

Tensor

Examples

>>> import torch
>>> def fn(x, y, a, b):
...    x1, x2 = x
...    y1, y2 = y
...    return (a*(x1 + x2) + b*(x1**2 + x1*x2 + x2**2))*(1 + y1 + y2)
>>> x = torch.tensor([0.0, 0.0])
>>> y = torch.zeros_like(x)
>>> s = series((2, 1), fn, x, y, 1.0, 1.0, retain=False)
>>> dx = torch.tensor([1.0, 2.0])
>>> dy = torch.tensor([3.0, 4.0])
>>> fn(x + dx, y + dy, 1.0, 1.0)
tensor(80.)
>>> evaluate(s, [dx, dy])
tensor(80.)

>>> from pprint import pprint
>>> import torch
>>> def fn(x, y, a, b):
...    x1, x2 = x
...    y1, y2 = y
...    return (a*(x1 + x2) + b*(x1**2 + x1*x2 + x2**2))*(1 + y1 + y2)
>>> x = torch.tensor([0.0, 0.0])
>>> y = torch.zeros_like(x)
>>> s = series((2, 1), fn, x, y, 1.0, 1.0, retain=False)
>>> s = series((2, 1), lambda x, y: evaluate(s, [x, y]), x, y, retain=False)
>>> pprint(s, sort_dicts=False)
{(0, 0, 0, 0): tensor(0.),
 (0, 0, 1, 0): tensor(0.),
 (0, 0, 0, 1): tensor(0.),
 (1, 0, 0, 0): tensor(1.),
 (0, 1, 0, 0): tensor(1.),
 (1, 0, 1, 0): tensor(1.),
 (1, 0, 0, 1): tensor(1.),
 (0, 1, 1, 0): tensor(1.),
 (0, 1, 0, 1): tensor(1.),
 (2, 0, 0, 0): tensor(1.),
 (1, 1, 0, 0): tensor(1.),
 (0, 2, 0, 0): tensor(1.),
 (2, 0, 1, 0): tensor(1.),
 (2, 0, 0, 1): tensor(1.),
 (1, 1, 1, 0): tensor(1.),
 (1, 1, 0, 1): tensor(1.),
 (0, 2, 1, 0): tensor(1.),
 (0, 2, 0, 1): tensor(1.)}

ndmap.gradient.factor(arrays: list[tuple[int, ...]]) → list[float]

Compute monomian factors given list of exponents

Parameters:

arrays (list[tuple[int, ...]], non-negative) – input indices

Return type:

list[float]

Examples

>>> factor([(1, 0)])
[1.0]
>>> factor([(2, 0)])
[0.5]
>>> factor([(2, 2)])
[0.25]
>>> factor([(2, 2, 2)])
[0.125]

ndmap.gradient.group(arrays: list[tuple[int, ...]], dimension: tuple[int, ...]) → list[tuple[int, ...]]

Standard indices grouping

Parameters:

• arrays (list[tuple[int, ...]], non-negative) – input indices

• dimension (tuple[int, ...], positive) – input dimension

Return type:

list[tuple[int, …]]

Examples

>>> xs = [(0, 2), (1, 1), (2, 0)]
>>> group(xs, (2, ))
[(2, 0), (1, 1), (0, 2)]
>>> xs = [(0, 2, 0, 1), (0, 2, 1, 0), (1, 1, 0, 1), (1, 1, 1, 0), (2, 0, 0, 1), (2, 0, 1, 0)]
>>> group(xs, (2, 2))
[(2, 0, 1, 0), (2, 0, 0, 1), (1, 1, 1, 0), (1, 1, 0, 1), (0, 2, 1, 0), (0, 2, 0, 1)]

ndmap.gradient.hessian(function: Callable) → Callable

Compute function hessian

Parameters:

function (Callable) – function

Return type:

Callable

Examples

>>> import torch
>>> def fn(x):
...    x1, x2, x3, x4 = x
...    return 1.0*x1**2 + 2.0*x2**2+ 3.0*x3**2 + 4.0*x4**2
>>> x = torch.tensor([0.0, 0.0, 0.0, 0.0])
>>> torch.func.hessian(fn)(x)
tensor([[2., 0., 0., 0.],
        [0., 4., 0., 0.],
        [0., 0., 6., 0.],
        [0., 0., 0., 8.]])
>>> hessian(fn)(x).detach().permute(1, 0)
tensor([[2., 0., 0., 0.],
        [0., 4., 0., 0.],
        [0., 0., 6., 0.],
        [0., 0., 0., 8.]])

ndmap.gradient.index(dimension: tuple[int, ...], order: tuple[int, ...], *, group: Callable | None = None, signature: list[tuple[int, ...]] | None = None) → list[tuple[int, ...]]

Generate monomial index table for a given dimension and order

Parameters:

• dimension (tuple[int, ...], positive) – monomial dimensions

• order (tuple[int, ...], non-negative) – derivative orders (total monomial degrees)

• group (Optional[Callable]) – grouping function

• signature (Optional[list[tuple[int, ...]]]) – allowed signatures

Returns:

monomial index table

Return type:

list[tuple[int, …]]

Examples

>>> index((2, ), (2, ))
[(0, 2), (1, 1), (2, 0)]
>>> index((4, ), (3, )) == index((2, 2), (2, 1))
True
>>> index((2, 2), (2, 1), signature=signature((2, 1)))
[(0, 2, 0, 1), (0, 2, 1, 0), (1, 1, 0, 1), (1, 1, 1, 0), (2, 0, 0, 1), (2, 0, 1, 0)]
>>> index((2, 2), (2, 1), signature=signature((2, 1)), group=group)
[(2, 0, 1, 0), (2, 0, 0, 1), (1, 1, 1, 0), (1, 1, 0, 1), (0, 2, 1, 0), (0, 2, 0, 1)]

ndmap.gradient.jacobian(function: Callable) → Callable

Compute function jacobian (can be composed)

Note, the output shape is different from jacfwd or jacrev

Parameters:

function (Callable) – function

Return type:

Callable

Examples

>>> import torch
>>> def fn(x):
...    x1, x2 = x
...    return torch.stack([1.0*x1 + 2.0*x2, 3.0*x1 + 4.0*x2])
>>> x = torch.tensor([0.0, 0.0])
>>> torch.func.jacrev(fn)(x)
tensor([[1., 2.],
        [3., 4.]])
>>> jacobian(fn)(x).detach().permute(1, 0)
tensor([[1., 2.],
        [3., 4.]])

>>> import torch
>>> def fn(x):
...    x1, x2 = x
...    y1 = torch.stack([1.0*x1 + 2.0*x2, 3.0*x1 + 4.0*x2])
...    y2 = torch.stack([5.0*x1 + 6.0*x2, 7.0*x1 + 8.0*x2])
...    return torch.stack([y1, y2])
>>> x = torch.tensor([0.0, 0.0])
>>> torch.func.jacrev(fn)(x).tolist()
[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]
>>> jacobian(fn)(x).detach().tolist()
[[[1.0, 3.0], [5.0, 7.0]], [[2.0, 4.0], [6.0, 8.0]]]
>>> jacobian(fn)(x).detach().permute(1, -1, 0).tolist()
[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]

ndmap.gradient.naught(state: torch.Tensor) → torch.Tensor

Infinitely differentiable zero function

Parameters:

state (Tensor) – input tensor

Return type:

Tensor

Examples

>>> import torch
>>> from torch.autograd import grad
>>> x = torch.tensor(0.0, requires_grad=True)
>>> y = (lambda x: x)(x)
>>> y, *_ = grad(y, x, retain_graph=True, create_graph=True)
>>> y
tensor(1.)
>>> y = (lambda x: x + naught(x))(x)
>>> y, *_ = grad(y, x, retain_graph=True, create_graph=True)
>>> y
tensor(1., grad_fn=<AddBackward0>)
>>> y, *_ = grad(y, x)
>>> y
tensor(-0.)

ndmap.gradient.reduce(array: tuple[int, ...]) → list[tuple[int, ...]]

Generate direct path from given index to zero

Parameters:

array (tuple[int, ...], non-negative) – input index

Return type:

list[tuple[int, …]]

Examples

>>> reduce((4, 0))
[(4, 0), (3, 0), (2, 0), (1, 0), (0, 0)]
>>> reduce((2, 2))
[(2, 2), (1, 2), (0, 2), (0, 1), (0, 0)]
>>> reduce((2, 2, 2))
[(2, 2, 2), (1, 2, 2), (0, 2, 2), (0, 1, 2), (0, 0, 2), (0, 0, 1), (0, 0, 0)]

ndmap.gradient.scalar(function: Callable, table: tuple[int]) → Callable

Given f(x, y, …) and table = map(len, (x, y, …)) return g(*x, *y, …) = f(x, y, …)

Parameters:

• function (Callable) – input function

• table (tuple[int, ...]) – map(len, (x, y, …))

• *pars (tuple) – passed to input function

Return type:

Callable

Examples

>>> import torch
>>> def fn(x, y):
...    x1, x2 = x
...    y1, y2, y3 = y
...    return x1*x2*y1*y2*y3
>>> def gn(x1, x2, y1, y2, y3):
...    return fn((x1, x2), (y1, y2, y3))
>>> x = torch.tensor([1, 1])
>>> y = torch.tensor([1, 1, 1])
>>> gn(*x, *y) == scalar(fn, (2, 3))(*x, *y)
tensor(True)

ndmap.gradient.select(function: ~typing.Callable, index: int, *, naught: ~typing.Callable = <function naught>) → Callable

Generate scalar function

Parameters:

• function (Callable) – function

• index (int, non-negative) – index

• naught (Callable, default=naught) – zero function

Return type:

Callable

Examples

>>> import torch
>>> def fn(x1, x2, x3, x4):
...    return torch.stack([x1, x2, x3, x4])
>>> x = torch.tensor([1.0, 2.0, 3.0, 4.0])
>>> [select(fn, i)(*x) for i in range(4)]
[tensor(1.), tensor(2.), tensor(3.), tensor(4.)]

Note

Input fuction is assumed to have scalar tensor arguments

ndmap.gradient.series(order: tuple[int, ...], function: ~typing.Callable, *args: tuple, retain: bool = True, series: bool = True, intermediate: bool | tuple[int, ...] = True, group: ~typing.Callable = <function group>, naught: ~typing.Callable = <function naught>, shape: tuple[int, ...] | None = None) → dict[tuple[int, ...], torch.Tensor]

Generate series representation of a given input function

c(i, j, k, …) * x**i * y**j * z**k * … => {…, (i, j, k, …) : c(i, j, k, …), …}

Note, the input function (returns a tensor) arguments are expected to be vector tensors

Parameters:

• order (tuple[int, ...], non-negative) – maximum derivative orders

• function (Callable) – input function

• *args (tuple) – input function arguments

• retain (bool, default=True) – flag to retain computation graph

• series (bool, default=True) – flag to return series coefficiens

• intermediate (Union[bool, tuple[int, ...]]) – flag to return indermidiate derivatives/coefficients

• group (Callable, default=group) – indices grouping function

• naught (Callable) – zero function

• shape (Optional[tuple[int, ...]]) – input function output shape

Return type:

dict[tuple[int, …], Tensor]

Examples

>>> from pprint import pprint
>>> import torch
>>> def fn(x, y, a, b):
...    x1, x2 = x
...    y1, y2 = y
...    return (a*(x1 + x2) + b*(x1**2 + x1*x2 + x2**2))*(1 + y1 + y2)
>>> x = torch.tensor([0.0, 0.0])
>>> y = torch.zeros_like(x)
>>> t = series((2, 1), fn, x, y, 1.0, 1.0, retain=False, series=True, intermediate=True)
>>> pprint(t, sort_dicts=False)
{(0, 0, 0, 0): tensor(0.),
 (0, 0, 1, 0): tensor(0.),
 (0, 0, 0, 1): tensor(0.),
 (1, 0, 0, 0): tensor(1.),
 (0, 1, 0, 0): tensor(1.),
 (1, 0, 1, 0): tensor(1.),
 (1, 0, 0, 1): tensor(1.),
 (0, 1, 1, 0): tensor(1.),
 (0, 1, 0, 1): tensor(1.),
 (2, 0, 0, 0): tensor(1.),
 (1, 1, 0, 0): tensor(1.),
 (0, 2, 0, 0): tensor(1.),
 (2, 0, 1, 0): tensor(1.),
 (2, 0, 0, 1): tensor(1.),
 (1, 1, 1, 0): tensor(1.),
 (1, 1, 0, 1): tensor(1.),
 (0, 2, 1, 0): tensor(1.),
 (0, 2, 0, 1): tensor(1.)}
>>> t = series((2, 1), fn, x, y, 1.0, 1.0, retain=False, series=True, intermediate=False)
>>> pprint(t, sort_dicts=False)
{(2, 0, 1, 0): tensor(1.),
 (2, 0, 0, 1): tensor(1.),
 (1, 1, 1, 0): tensor(1.),
 (1, 1, 0, 1): tensor(1.),
 (0, 2, 1, 0): tensor(1.),
 (0, 2, 0, 1): tensor(1.)}
>>> t = series((2, 1), fn, x, y, 1.0, 1.0, retain=False, series=True, intermediate=(2, 0, 1, 0))
>>> pprint(t, sort_dicts=False)
{(2, 0, 1, 0): tensor(1.)}
>>> t = series((2, 1), fn, x, y, 1.0, 1.0, retain=False, series=False, intermediate=(2, 0, 1, 0))
>>> pprint(t, sort_dicts=False)
{(2, 0, 1, 0): tensor(2.)}
>>> t = series((2, 1), fn, x, y, 1.0, 1.0, retain=True, series=True, intermediate=(2, 0, 1, 0))
>>> pprint(t, sort_dicts=False)
{(2, 0, 1, 0): tensor(1., grad_fn=<MulBackward0>)}

ndmap.gradient.signature(order: tuple[int, ...]) → list[tuple[int, ...]]

Compute derivative signatures from given total group orders

Parameters:

order (tuple[int, ...], non-negative) – tuple of orders

Returns:

list of signatures

Return type:

list[tuple[int, …]]

Examples

>>> signature((2, ))
[(0,), (1,), (2,)]
>>> signature((1, 2))
[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]
>>> signature((2, 1))
[(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)]
>>> signature((2, 2))
[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)]

ndmap.gradient.split(array: tuple[int, ...], chunks: tuple[int, ...]) → list[tuple[int, ...]]

Split array into chuncks with given length

Parameters:

• array (tuple[int, ...]) – array to split

• chunks (tuple[int, ...], positive) – list of chunks to use

Return type:

list[tuple[int, …]]

Examples

>>> split((1, 2, 3, 4, 5), (2, 3))
[(1, 2), (3, 4, 5)]
>>> split((1, 2, 3, 4, 5), (2, 1))
[(1, 2), (3,)]
>>> split((1, 2, 3, 4, 5), (2, 1, 2))
[(1, 2), (3,), (4, 5)]