# Copyright (c) OpenMMLab. All rights reserved.
import torch

from juxtapose.mmdeploy.core import FUNCTION_REWRITER


@FUNCTION_REWRITER.register_rewriter(func_name="torch.triu")
@FUNCTION_REWRITER.register_rewriter(func_name="torch.Tensor.triu")
def triu__default(
    input: torch.Tensor, diagonal: int = 0, *args, **kwargs
) -> torch.Tensor:
    """Rewrite `triu` for exporting model to ONNX."""
    assert len(input.shape) >= 2
    height, width = input.shape[-2:]
    arange0 = torch.arange(width, device=input.device).unsqueeze(0)
    arange1 = torch.arange(height, device=input.device).unsqueeze(-1)
    mask = arange0 >= torch.add(arange1, diagonal)
    return input * mask
