Home
last modified time | relevance | path

Searched refs:compute_local_shape_and_global_offset (Results 1 – 6 of 6) sorted by relevance

/aosp_15_r20/external/pytorch/test/distributed/_tensor/
H A Dtest_utils.py9 compute_local_shape_and_global_offset,
76 local_size, global_offset = compute_local_shape_and_global_offset(
108 local_size, global_offset = compute_local_shape_and_global_offset(
137 local_shape, global_offset = compute_local_shape_and_global_offset(
164 local_shape, global_offset = compute_local_shape_and_global_offset(
200 _, _ = compute_local_shape_and_global_offset(
221 _, _ = compute_local_shape_and_global_offset(
233 _, _ = compute_local_shape_and_global_offset(
255 local_shape, _ = compute_local_shape_and_global_offset(
273 _, _ = compute_local_shape_and_global_offset(
[all …]
H A Dtest_random_ops.py10 from torch.distributed._tensor._utils import compute_local_shape_and_global_offset
244 _, local_shard_offset = compute_local_shape_and_global_offset(
H A Dtest_dtensor.py685 compute_local_shape_and_global_offset,
692 _, offset = compute_local_shape_and_global_offset(
/aosp_15_r20/external/pytorch/torch/distributed/checkpoint/
H A Dplanner_helpers.py11 from torch.distributed._tensor._utils import compute_local_shape_and_global_offset
75 sizes, offsets = compute_local_shape_and_global_offset(
241 sizes, offsets = compute_local_shape_and_global_offset(
/aosp_15_r20/external/pytorch/torch/distributed/
H A D_state_dict_utils.py31 from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
556 shape, offset = compute_local_shape_and_global_offset(
/aosp_15_r20/external/pytorch/torch/distributed/tensor/
H A D_utils.py49 def compute_local_shape_and_global_offset( function