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

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"""Sinkhorn normalization functions used by manifold hyper-connections.""" 

16 

17from typing import Any, Tuple 

18 

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

20 

21try: 

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

23except ImportError: 

24 omni_training_custom_ops = None 

25 

26 

27class _Sinkhorn(torch.autograd.Function): 

28 """Autograd bridge for the NPU Sinkhorn custom operator.""" 

29 

30 @staticmethod 

31 def forward( 

32 ctx: Any, 

33 h_res: torch.Tensor, 

34 sinkhorn_iters: int = 20, 

35 eps: float = 1e-6, 

36 ) -> torch.Tensor: 

37 """Run the NPU Sinkhorn forward operator.""" 

38 if omni_training_custom_ops is None: 

39 raise ImportError("NPU Sinkhorn requires omni_training_custom_ops") 

40 output, norm_out, sum_out = torch.ops.custom.npu_sinkhorn( 

41 h_res, 

42 out_flag=1, 

43 eps=eps, 

44 num_iters=sinkhorn_iters, 

45 ) 

46 ctx.save_for_backward(norm_out, sum_out) 

47 return output 

48 

49 @staticmethod 

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

51 """Run the NPU Sinkhorn backward operator.""" 

52 norm_out, sum_out = ctx.saved_tensors 

53 grad_h_res = torch.ops.custom.npu_sinkhorn_grad(grad_output, norm_out, sum_out) 

54 grads = [grad_h_res, None, None] 

55 return tuple(grads) 

56 

57 

58def sinkhorn(h_res: torch.Tensor, sinkhorn_iters: int = 20, eps: float = 1e-6) -> torch.Tensor: 

59 """Normalize a square mixing matrix with the NPU Sinkhorn operator. 

60 

61 Args: 

62 h_res: Residual mixing matrix. 

63 sinkhorn_iters: Number of normalization iterations. 

64 eps: Numerical stability epsilon. 

65 

66 Returns: 

67 The normalized residual mixing matrix. 

68 """ 

69 return _Sinkhorn.apply(h_res, sinkhorn_iters, eps) 

70 

71 

72def sinkhorn_knopps(h_res: torch.Tensor, sinkhorn_iters: int = 20, eps: float = 1e-6) -> torch.Tensor: 

73 """Normalize a square mixing matrix with the portable PyTorch algorithm. 

74 

75 Args: 

76 h_res: Residual mixing logits. 

77 sinkhorn_iters: Number of normalization iterations. 

78 eps: Numerical stability epsilon. 

79 

80 Returns: 

81 The normalized residual mixing matrix. 

82 """ 

83 h_res = h_res.softmax(-1) + eps 

84 col_sum = h_res.sum(-2, keepdim=True) 

85 h_res = h_res / (col_sum + eps) 

86 for _ in range(sinkhorn_iters - 1): 

87 row_sum = h_res.sum(-1, keepdim=True) 

88 h_res = h_res / (row_sum + eps) 

89 col_sum = h_res.sum(-2, keepdim=True) 

90 h_res = h_res / (col_sum + eps) 

91 return h_res