treex.Map

Mixin that adds a .map() method to the class.

Source code in treeo/mixins.py

class Map:
    """
    Mixin that adds a `.map()` method to the class.
    """

    def map(
        self: A,
        f: tp.Callable,
        *filters: api.Filter,
        inplace: bool = False,
        flatten_mode: tp.Union[api.FlattenMode, str, None] = None,
        is_leaf: tp.Callable[[tp.Any], bool] = None,
    ) -> A:
        """
        `map` is a wrapper over `treeo.map` that passes `self` as the second argument.

        Arguments:
            f: The function to apply to the leaves.
            *filters: The filters used to select the leaves to which the function will be applied.
            inplace: If `True`, the input `obj` is mutated and returned.
            flatten_mode: Sets a new `FlattenMode` context for the operation, if `None` the current context is used.

        Returns:
            A new pytree with the changes applied. If `inplace` is `True`, the input `obj` is returned.
        """
        return api.map(
            f,
            self,
            *filters,
            inplace=inplace,
            flatten_mode=flatten_mode,
            is_leaf=is_leaf,
        )

map(self, f, *filters, *, inplace=False, flatten_mode=None, is_leaf=None)

map is a wrapper over treeo.map that passes self as the second argument.

Parameters:

Name	Type	Description	Default
f	Callable	The function to apply to the leaves.	required
*filters	Union[Type[Any], Callable[[FieldInfo], bool]]	The filters used to select the leaves to which the function will be applied.	()
inplace	bool	If True, the input obj is mutated and returned.	False
flatten_mode	Union[treeo.tree.FlattenMode, str]	Sets a new FlattenMode context for the operation, if None the current context is used.	None

Returns:

Type	Description
~A	A new pytree with the changes applied. If inplace is True, the input obj is returned.

Source code in treeo/mixins.py

def map(
    self: A,
    f: tp.Callable,
    *filters: api.Filter,
    inplace: bool = False,
    flatten_mode: tp.Union[api.FlattenMode, str, None] = None,
    is_leaf: tp.Callable[[tp.Any], bool] = None,
) -> A:
    """
    `map` is a wrapper over `treeo.map` that passes `self` as the second argument.

    Arguments:
        f: The function to apply to the leaves.
        *filters: The filters used to select the leaves to which the function will be applied.
        inplace: If `True`, the input `obj` is mutated and returned.
        flatten_mode: Sets a new `FlattenMode` context for the operation, if `None` the current context is used.

    Returns:
        A new pytree with the changes applied. If `inplace` is `True`, the input `obj` is returned.
    """
    return api.map(
        f,
        self,
        *filters,
        inplace=inplace,
        flatten_mode=flatten_mode,
        is_leaf=is_leaf,
    )