Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / aggregate_hidden.py: 0%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Masked hidden-state aggregation function.""" 

16 

17from typing import Any, Tuple 

18 

19import torch # pylint: disable=forbidden-backend-import 

20import omni_training_custom_ops # noqa: F401 # pylint: disable=unused-import 

21 

22 

23class _AggregateHidden(torch.autograd.Function): 

24 """Autograd bridge for the aggregate-hidden custom operator.""" 

25 

26 @staticmethod 

27 def forward( 

28 ctx: Any, 

29 input_tensor: torch.Tensor, 

30 weight: torch.Tensor, 

31 mask: torch.Tensor, 

32 ) -> torch.Tensor: 

33 """Run the aggregate-hidden forward operator.""" 

34 ctx.save_for_backward(input_tensor, weight, mask) 

35 return torch.ops.custom.npu_aggregate_hidden(input_tensor, weight, mask=mask) 

36 

37 @staticmethod 

38 def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, None]: 

39 """Run the aggregate-hidden backward operator.""" 

40 input_tensor, weight, mask = ctx.saved_tensors 

41 grad_input, grad_weight = torch.ops.custom.npu_aggregate_hidden_grad( 

42 grad_output.contiguous(), 

43 input_tensor, 

44 weight, 

45 mask=mask, 

46 ) 

47 return grad_input, grad_weight, None 

48 

49 

50def aggregate_hidden(input_tensor: torch.Tensor, weight: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: 

51 """Aggregate masked hidden states with the NPU custom operator. 

52 

53 Args: 

54 input_tensor: Hidden states in ``[sequence, batch, hidden]`` layout. 

55 weight: Depthwise aggregation weights in ``[window, hidden]`` layout. 

56 mask: Boolean validity mask in ``[batch, sequence]`` layout. 

57 

58 Returns: 

59 Aggregated hidden states. 

60 """ 

61 return _AggregateHidden.apply(input_tensor, weight, mask)