Function _map_structure

dali/python/nvidia/dali/_conditionals.py:53–58

Custom wrapper over optree.tree_map that filters it out from the user-visible stack trace for error reporting purposes.

(func, *structures, **kwargs)

def _map_structure(func, *structures, **kwargs):
    """Custom wrapper over optree.tree_map that filters it out from the user-visible stack trace
    for error reporting purposes.
    """
    with _autograph.CustomModuleFilter(optree):
        return optree.tree_map(func, *structures, **kwargs)


class _Branch(Enum):