liblaf.jarp.tree ¶

def array(**kwargs: Unpack[FieldOptions[Any]]) -> Array:
    """Create a data field whose default is normalized to a JAX array.

    When `default` is a concrete array-like value, `array` rewrites it into
    a factory so each instance receives its own array object.
    """
    if "default" in kwargs and "factory" not in kwargs:
        default: Any = kwargs["default"]
        if not (default is None or isinstance(default, attrs.Factory)):  # ty:ignore[invalid-argument-type]
            default: Array = jnp.asarray(default)
            kwargs.pop("default")
            kwargs["factory"] = lambda: default
    return field(**kwargs)  # ty:ignore[no-matching-overload]

@_wraps(attrs.field)
def auto(**kwargs) -> Any:
    """Create a field whose PyTree role is chosen from the runtime value."""
    kwargs.setdefault("static", FieldType.AUTO)
    return field(**kwargs)

def codegen_pytree_functions(
    cls: type,
    data_fields: Sequence[str] = (),
    meta_fields: Sequence[str] = (),
    auto_fields: Sequence[str] = (),
    *,
    filter_spec: Callable[[Any], bool] = is_data,
    bypass_setattr: bool | None = None,
) -> PyTreeFunctions:
    """Generate flatten and unflatten callbacks for a class.

    Args:
        cls: Class whose instances should become PyTree nodes.
        data_fields: Field names that are always emitted as dynamic children.
        meta_fields: Field names that are always emitted as auxiliary metadata.
        auto_fields: Field names filtered at runtime with `filter_spec`.
        filter_spec: Predicate used to split `auto_fields` into dynamic data
            or metadata.
        bypass_setattr: Whether generated unflattening code should use
            [`object.__setattr__`][object.__setattr__] instead of normal
            attribute assignment.

    Returns:
        A [`PyTreeFunctions`][liblaf.jarp.tree.codegen.PyTreeFunctions] tuple
        containing `flatten`, `unflatten`, and `flatten_with_keys` callables.
    """
    if bypass_setattr is None:
        bypass_setattr = cls.__setattr__ is not object.__setattr__
    flatten_def: ast.FunctionDef = codegen_flatten(
        data_fields, meta_fields, auto_fields
    )
    flatten_with_keys_def: ast.FunctionDef = codegen_flatten_with_keys(
        data_fields, meta_fields, auto_fields
    )
    unflatten_def: ast.FunctionDef = codegen_unflatten(
        data_fields, meta_fields, auto_fields, bypass_setattr=bypass_setattr
    )
    module: ast.Module = ast.Module(
        body=[flatten_def, flatten_with_keys_def, unflatten_def], type_ignores=[]
    )
    module = ast.fix_missing_locations(module)
    source: str = ast.unparse(module)
    namespace: dict = {
        "_cls": cls,
        "_filter_spec": filter_spec,
        "_object_new": object.__new__,
        "_object_setattr": object.__setattr__,
        **_make_keys((*data_fields, *meta_fields, *auto_fields)),
    }
    filename: str = _make_filename(cls)
    # use unparse source so we have correct source code locations
    code: types.CodeType = compile(source, filename, "exec")
    exec(code, namespace)  # noqa: S102
    _update_linecache(source, filename)
    return PyTreeFunctions(
        _add_dunder(cls, namespace["flatten"]),
        _add_dunder(cls, namespace["unflatten"]),
        _add_dunder(cls, namespace["flatten_with_keys"]),
    )

def combine[T](data_leaves: Iterable[Array | None], aux: AuxData[T]) -> T:
    """Rebuild a PyTree from dynamic leaves and recorded metadata."""
    leaves: list[Any] = combine_leaves(data_leaves, aux.meta_leaves)
    return jax.tree.unflatten(aux.treedef, leaves)

def combine_leaves(
    data_leaves: Iterable[Array | None], meta_leaves: Iterable[Any]
) -> list[Any]:
    """Merge dynamic leaves back together with their static counterparts."""
    return [
        data_leaf if meta_leaf is None else meta_leaf
        for data_leaf, meta_leaf in zip(data_leaves, meta_leaves, strict=True)
    ]

def define[T: type](maybe_cls: T | None = None, **kwargs: Any) -> Any:
    """Define an `attrs` class and optionally register it as a PyTree.

    Args:
        maybe_cls: Class being decorated. When omitted, return a configured
            decorator.
        **kwargs: Options forwarded to [`attrs.define`][attrs.define], plus
            `pytree` to control JAX registration. `pytree="data"`
            registers fields with `fieldz` semantics, `"static"` registers
            the whole instance as a static value, and `"none"` leaves the
            class unregistered.

    Returns:
        The decorated class or a class decorator.
    """
    if maybe_cls is None:
        return functools.partial(define, **kwargs)
    pytree: PyTreeType = PyTreeType(kwargs.pop("pytree", None))
    frozen: bool = kwargs.get("frozen", False)
    if pytree is PyTreeType.STATIC and not frozen:
        warnings.warn(
            "Defining a static class that is not frozen may lead to unexpected behavior.",
            stacklevel=2,
        )
    cls: T = attrs.define(maybe_cls, **kwargs)  # ty:ignore[invalid-assignment]
    match pytree:
        case PyTreeType.DATA:
            register_fieldz(cls)
        case PyTreeType.STATIC:
            jtu.register_static(cls)
    return cls

@_wraps(attrs.field)
def field(**kwargs) -> Any:
    """Create an `attrs` field using jarp's `static` metadata convention."""
    if "static" in kwargs:
        kwargs["metadata"] = {
            "static": kwargs.pop("static"),
            **(kwargs.get("metadata") or {}),
        }
    return attrs.field(**kwargs)

def frozen[T: type](maybe_cls: T | None = None, **kwargs: Any) -> Any:
    """Define a frozen `attrs` class and register it as a data PyTree.

    This is the common choice for immutable structures whose array fields
    should participate in JAX transformations.
    """
    _warnings_hide = True
    if maybe_cls is None:
        return functools.partial(frozen, **kwargs)
    kwargs.setdefault("frozen", True)
    return define(maybe_cls, **kwargs)

def frozen_static[T: type](maybe_cls: T | None = None, **kwargs: Any) -> Any:
    """Define a frozen `attrs` class and register it as a static PyTree.

    Use this for immutable helper objects that should be treated as static
    metadata instead of flattening into JAX leaves.
    """
    _warnings_hide = True
    if maybe_cls is None:
        return functools.partial(frozen_static, **kwargs)
    kwargs.setdefault("frozen", True)
    kwargs.setdefault("pytree", PyTreeType.STATIC)
    return define(maybe_cls, **kwargs)

def is_data(obj: Any) -> bool:
    """Return whether a value stays on the dynamic side of a partition.

    Dynamic values include JAX arrays, `None` placeholders, and objects whose
    type already has a JAX PyTree registration. Everything else is treated as
    static metadata by [`partition`][liblaf.jarp.tree.partition].
    """
    return is_leaf(obj) or jtu.is_tree_node(type(obj))

def is_leaf(obj: Any) -> TypeIs[Array | None]:
    """Return whether a leaf contributes numeric data to a flat vector.

    This is intentionally narrower than [`is_data`][liblaf.jarp.tree.is_data]:
    only arrays and `None` participate in the flat-vector protocol used by
    [liblaf.jarp.ravel][].
    """
    return obj is None or isinstance(obj, Array)

def partial[T](func: Callable[..., T], /, *args: Any, **kwargs: Any) -> Partial[..., T]:
    """Partially apply a callable and keep bound values visible to JAX trees."""
    return Partial(func, *args, **kwargs)

def partition[T](obj: T) -> tuple[list[Array | None], AuxData[T]]:
    """Split a PyTree into dynamic leaves and static metadata.

    The returned leaf list preserves tree order. Non-dynamic positions become
    `None` in the data list and are stored in the accompanying
    [`AuxData`][liblaf.jarp.tree.AuxData].
    """
    leaves, treedef = jax.tree.flatten(obj)
    data_leaves, meta_leaves = partition_leaves(leaves)
    return data_leaves, AuxData(tuple(meta_leaves), treedef)

def partition_leaves(leaves: list[Any]) -> tuple[list[Array | None], list[Any]]:
    """Separate raw tree leaves into dynamic leaves and metadata leaves."""
    data_leaves: list[Array | None] = []
    meta_leaves: list[Any] = []
    for leaf in leaves:
        if is_leaf(leaf):
            data_leaves.append(leaf)
            meta_leaves.append(None)
        else:
            data_leaves.append(None)
            meta_leaves.append(leaf)
    return data_leaves, meta_leaves

def ravel[T](tree: T) -> tuple[Array, Structure[T]]:
    """Flatten a PyTree's dynamic leaves into one vector.

    Non-array leaves are treated as static metadata and preserved in the
    returned [`Structure`][liblaf.jarp.tree.Structure] instead of being
    concatenated into the flat array.

    Examples:
        >>> import jax.numpy as jnp
        >>> from liblaf import jarp
        >>> flat, structure = jarp.ravel({"x": jnp.array([1.0, 2.0]), "tag": "train"})
        >>> flat.tolist()
        [1.0, 2.0]
        >>> rebuilt = structure.unravel(jnp.array([3.0, 4.0]))
        >>> rebuilt["x"].tolist(), rebuilt["tag"]
        ([3.0, 4.0], 'train')

    Args:
        tree: PyTree to flatten.

    Returns:
        A tuple of `(flat, structure)` where `flat` is a one-dimensional JAX
        array and `structure` can rebuild compatible trees later.
    """
    leaves, treedef = jax.tree.flatten(tree)
    dynamic_leaves, static_leaves = partition_leaves(leaves)
    flat: Array = _ravel(dynamic_leaves)
    structure: Structure[T] = Structure(
        offsets=_offsets_from_leaves(dynamic_leaves),
        shapes=_shapes_from_leaves(dynamic_leaves),
        meta_leaves=tuple(static_leaves),
        treedef=treedef,
        dtype=flat.dtype,
    )
    return flat, structure

def register_fieldz[T: type](
    cls: T,
    data_fields: Sequence[str] | None = None,
    meta_fields: Sequence[str] | None = None,
    auto_fields: Sequence[str] | None = None,
    *,
    filter_spec: Callable[[Any], bool] = is_data,
    bypass_setattr: bool | None = None,
) -> T:
    """Register an `attrs` class with JAX using field metadata.

    Field groups default to the metadata written by
    [`array`][liblaf.jarp.tree.array], [`auto`][liblaf.jarp.tree.auto], and
    [`static`][liblaf.jarp.tree.static]. Pass explicit field lists when you
    need to register a class that was not declared with `liblaf.jarp` field
    helpers.

    Args:
        cls: Class to register.
        data_fields: Field names that are always treated as dynamic children.
        meta_fields: Field names that are always treated as static metadata.
        auto_fields: Field names filtered at runtime with `filter_spec`.
        filter_spec: Predicate used to split `auto_fields` into dynamic data
            or metadata.
        bypass_setattr: Whether generated unflattening code should use
            [`object.__setattr__`][object.__setattr__] instead of normal
            attribute assignment.

    Returns:
        The same class object, for decorator-style usage.
    """
    if data_fields is None:
        data_fields: list[str] = _filter_field_names(cls, FieldType.DATA)
    if meta_fields is None:
        meta_fields: list[str] = _filter_field_names(cls, FieldType.META)
    if auto_fields is None:
        auto_fields: list[str] = _filter_field_names(cls, FieldType.AUTO)
    register_generic(
        cls,
        data_fields,
        meta_fields,
        auto_fields,
        filter_spec=filter_spec,
        bypass_setattr=bypass_setattr,
    )
    return cls

def register_generic(
    cls: type,
    data_fields: Sequence[str] = (),
    meta_fields: Sequence[str] = (),
    auto_fields: Sequence[str] = (),
    *,
    filter_spec: Callable[[Any], bool] = is_data,
    bypass_setattr: bool | None = None,
) -> None:
    """Register a class as a PyTree using explicit field groups.

    Use this lower-level helper when you want to control the flattening layout
    directly instead of relying on [attrs][] metadata.

    Args:
        cls: Class to register.
        data_fields: Field names that are always emitted as dynamic children.
        meta_fields: Field names that are always emitted as auxiliary metadata.
        auto_fields: Field names filtered at runtime with `filter_spec`.
        filter_spec: Predicate used to split `auto_fields` into dynamic data
            or metadata.
        bypass_setattr: Whether generated unflattening code should use
            [`object.__setattr__`][object.__setattr__] instead of normal
            attribute assignment.
    """
    flatten: Callable
    unflatten: Callable
    flatten_with_keys: Callable
    flatten, unflatten, flatten_with_keys = codegen_pytree_functions(
        cls,
        data_fields,
        meta_fields,
        auto_fields,
        filter_spec=filter_spec,
        bypass_setattr=bypass_setattr,
    )
    jtu.register_pytree_node(cls, flatten, unflatten, flatten_with_keys)

def select[T](
    condlist: Sequence[Bool[ArrayLike, " ..."]], choicelist: Sequence[T], default: T
) -> T:
    """Select among matching PyTrees with `jax.numpy.select`.

    Each leaf is selected independently with ordered conditions. The first true
    condition at each position selects the corresponding choice leaf; `default`
    supplies the leaf where no condition is true.

    Args:
        condlist: Non-empty sequence of boolean scalar or array-like
            conditions.
        choicelist: PyTrees to choose from. It must have the same length as
            `condlist`, and every choice must have the same tree structure as
            `default`.
        default: PyTree returned where no condition is true.

    Returns:
        A PyTree with the same structure as `default`.

    Raises:
        ValueError: If `condlist` is empty or its length does not match
            `choicelist`.

    Examples:
        >>> import jax.numpy as jnp
        >>> from liblaf import jarp
        >>> result = jarp.tree.select(
        ...     [jnp.array([False, True, False]), jnp.array([True, True, False])],
        ...     [{"value": jnp.array([1, 1, 1])}, {"value": jnp.array([2, 2, 2])}],
        ...     {"value": jnp.array([9, 9, 9])},
        ... )
        >>> result["value"].tolist()
        [2, 1, 9]
    """
    return jax.tree.map(
        lambda *args: jnp.select(condlist, args[:-1], args[-1]), *choicelist, default
    )

@_wraps(attrs.field)
def static(**kwargs) -> Any:
    """Create a field that is always treated as static metadata."""
    # for consistency with `jax.tree_util.register_dataclass`
    kwargs.setdefault("static", True)
    return field(**kwargs)

def where[T](condition: Bool[ArrayLike, " ..."], x: T, y: T) -> T:
    """Choose between matching PyTrees with `jax.numpy.where`.

    Args:
        condition: Boolean scalar or array-like condition.
        x: PyTree used where `condition` is true.
        y: PyTree used where `condition` is false.

    Returns:
        A PyTree with the same structure as `x` and `y`.

    Examples:
        >>> import jax.numpy as jnp
        >>> from liblaf import jarp
        >>> result = jarp.tree.where(
        ...     jnp.array([True, False]),
        ...     {"value": jnp.array([1, 2])},
        ...     {"value": jnp.array([3, 4])},
        ... )
        >>> result["value"].tolist()
        [1, 4]
    """
    return jax.tree.map(lambda a, b: jnp.where(condition, a, b), x, y)