torch-ext/selective_update/__init__.py

from typing import Optional

import torch

from .functional import selective_state_update


def selective_update(x: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
    """Legacy scaffold kernel — adds 1.0 to each element."""
    if out is None:
        out = torch.empty_like(x)
    out.copy_(x + 1.0)
    return out


__all__ = ["selective_update", "selective_state_update"]