| """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 | |