"""Project-local compatibility patches for the H20 CoVT environment."""

try:
    import torch
    import flash_attn.bert_padding as _flash_bert_padding

    def _covt_pad_input(hidden_states, indices, batch, seqlen):
        trailing_shape = tuple(hidden_states.shape[1:])
        output = torch.zeros(
            (batch * seqlen, *trailing_shape),
            device=hidden_states.device,
            dtype=hidden_states.dtype,
        )
        output[indices] = hidden_states
        return output.reshape(batch, seqlen, *trailing_shape)

    _flash_bert_padding.pad_input = _covt_pad_input
except Exception:
    pass