Function prune_tree

axlearn/common/utils.py:1427–1458

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

(
    in_tree: NestedTensor,
    should_prune: Callable[[str, NestedTensor], bool],
    *,
    prefix: str = "",
    separator: str = "/",
)

            default_dtype=param_dtype,
        )
    # Check if param_dtype is an instance of ConfigOr[PerParamFn[jnp.dtype]]
    elif isinstance(param_dtype, PerParamFn) or (
        isinstance(param_dtype, FunctionConfigBase) and isinstance(param_dtype.fn, PerParamFn)
    ):
        return param_dtype
    raise ValueError(
        f"{param_dtype} does not match any required types, should be "
        "jnp.dtype or ConfigOr[PerParamFn[jnp.dtype]]."
    )


def count_model_params(tree: NestedTensor) -> int:
    """Count the number of parameters in a model."""
    return sum(x.size for x in jax.tree_util.tree_leaves(tree))


def check_param_shape_alignment(
    source_tree: NestedTensor, target_tree: NestedTensor
) -> Union[None, str]:
    """Check param shape alignment between two parameter trees.

    This function assumes both trees have the same structures.

    Args:
        source_tree: The source parameter tree,
          which can be obtained via trainer.trainer_state.model or else.
        target_tree: The target parameter tree,
          which can be obtained via trainer.trainer_state.model or else.

    Returns:
        None if shape matches.
        A message indicating which parameter shapes are mismatched.
        e.g. "(linear1/weight/0) shape is different: source: (32), target: (15)."
    """
    param_shape_source = jax.tree.map(lambda x: x.shape, source_tree)