# 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,
)