Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / utils / shape_utils.py: 100%

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

16Utility functions for distributed tensor operations. 

17 

18This module provides helper functions for computing local shapes, global offsets, 

19and other layout-related calculations in distributed settings. 

20""" 

21from hyper_parallel.core.dtensor.layout import Layout 

22 

23 

24def compute_local_shape_and_global_offset(global_shape, device_mesh, placement): 

25 """ 

26 Compute local shard shape and its global offset. 

27 

28 Args: 

29 global_shape: Shape of the global tensor. 

30 mesh: Device mesh for distributed execution. 

31 placements: Sharding placements for each dimension. 

32 Supports Placement objects or alias strings. 

33 

34 Returns: 

35 tuple: (local_shape, global_offset) 

36 """ 

37 from hyper_parallel.core.dtensor.dtensor import _is_alias_placements # pylint: disable=C0415 

38 total_layout = Layout.from_device_mesh(device_mesh) 

39 if _is_alias_placements(placement): 

40 layout = total_layout(*placement) 

41 else: 

42 layout = total_layout(placement) 

43 layout.placement_to_tensor_map(len(global_shape)) 

44 slice_shape = list(global_shape) 

45 alias_tensor_map = layout.alias_tensor_map 

46 for i, axis_name in enumerate(alias_tensor_map): 

47 if isinstance(axis_name, str): 

48 axis_name = (axis_name,) 

49 for sub_axis_name in axis_name: 

50 if sub_axis_name != "None": 

51 num_devices = layout.mesh.get_device_num_along_axis(sub_axis_name) 

52 local_rank = layout.mesh.get_local_rank(sub_axis_name) 

53 global_size = slice_shape[i] 

54 remainder = global_size % num_devices 

55 # Consistent with torch.chunk: first `remainder` ranks get one extra element 

56 if remainder != 0 and local_rank < remainder: 

57 slice_shape[i] = global_size // num_devices + 1 

58 else: 

59 slice_shape[i] = global_size // num_devices 

60 return slice_shape