corewise_map

t3toolbox.corewise.corewise_map(f, *Xs)

def corewise_map(
        f,                  # callable applied to matching leaves: f(leaf_X0, leaf_X1, ...) -> leaf
        *Xs:  NDArrayTree,  # one or more identically-structured trees
) -> NDArrayTree:           # tree of f-results, same structure as the inputs

Apply f elementwise over the leaves of one or more identically-structured trees (the general tree map the corewise ops are special cases of; used e.g. for an Adam moment update over the cores).

Examples

>>> import numpy as np
>>> import t3toolbox.corewise as cw
>>> X = (np.ones(2), (3.0,))
>>> print(cw.corewise_map(lambda x: x + 1, X))
(array([2., 2.]), (4.0,))
>>> print(cw.corewise_map(lambda a, b: a * b, X, X))
(array([1., 1.]), (9.0,))

Parameters:

Xs (NDArrayTree)

Return type:

NDArrayTree