Function set_data_dir

axlearn/common/utils.py:1532–1556

Sets the environment variable DATA_DIR to the given `data_dir`. Args: data_dir: The data_dir. Raises: ValueError: If the environment variable DATA_DIR is already set to a different value.

(data_dir: Optional[str])

    partial_fn = functools.partial(fn, *args, **kwargs)
    return functools.update_wrapper(partial_fn, fn)


def prune_tree(
    in_tree: NestedTensor,
    should_prune: Callable[[str, NestedTensor], bool],
    *,
    prefix: str = "",
    separator: str = "/",
):
    """Returns a shallow copy of the input tree with subtrees pruned based on `should_prune`.

    This is a shallow copy because leaf nodes (non-dict values) are not deep-copied.

    Args:
        in_tree: The input tree to be pruned.
        should_prune: A callable which takes (path, subtree) as input and returns a boolean. The
            subtree provided will have already been pruned. If the callable returns True, the
            subtree itself will be dropped.
        prefix: Path prefix.
        separator: Separator used to join path parts.

    Returns:
        The pruned copy of the input tree.
    """
    if isinstance(in_tree, dict):
        # Use type() so that if in_tree is a VDict, out_tree is also a VDict.
        out_tree = type(in_tree)()
        for k, v in in_tree.items():