【合入来源】

https://gitcode.com/Ascend/pytorch-ecosystem/issues/326
请勿携带内部流程信息（需求链接、问题单、内部issue等）

• 需求
• 问题单
• issue/工单
• 重构优化
• 资料更新

【修改方案】

请描述修改内容的具体实现，涉及哪些组件之间进行交互，可以用1、2、3、...进行罗列
如果是需求或者重构类的PR，需要补充详细设计文档（说明上下游组件关系、时序图、类图、DFX能力等内容）

1. API 功能说明：torch.distributed.tensor._utils.compute_local_shape_and_global_offset 是 DTensor 模块下一个纯 Python 实现的内部元数据工具函数。其功能为：给定全局张量的形状、设备网格（DeviceMesh）以及放置策略（placements），计算并推导出当前 rank 对应的局部张量形状（local_shape）以及该数据块在全局张量中的绝对偏移量（global_offset）。
2. 新增测试用例的原因：经排查 PyTorch 原生测试代码，官方原有的 1D 和 2D 测试用例中仅覆盖了 Shard（切分）和 Replicate（复制）策略，客观上遗漏了对另一合法核心策略 Partial（部分状态）的入参验证。为了保证该 API 在 NPU 环境下的参数完备性与逻辑全覆盖，必须进行用例补齐。
3. 实现细节：新增 test_compute_local_shape_and_global_offset_partial 测试用例，显式导入 Partial 策略类，并断言验证其推导结果：在 Partial 策略下，局部张量尺寸未发生物理切分（保持等于全局尺寸），且全局偏移坐标全为 0。
4. 修改test_utils.py存在的bug：

具体的bug：

经过检查原第5行中的 torch.distributed._tensor._utils.compute_local_shape 已不存在，其已被替换成 torch.distributed._tensor._utils.compute_local_shape.compute_local_shape_and_global_offset 。

检查结果如下， compute_local_shape 已不存在:

torch.distributed._tensor._utils.compute_local_shape.compute_local_shape_and_global_offset函数返回一个元组 (local_shape, global_offset)，其中第一个元素 local_shape 与原 torch.distributed._tensor._utils.compute_local_shape 的返回值完全一致。因此通过元组解包 local_shape, _ = compute_local_shape_and_global_offset(...) 丢弃第二个返回值即可。

原来代码：

 local_size = compute_local_shape(size, mesh, placements)

修改为：

local_size, _ = compute_local_shape_and_global_offset(size, mesh, placements)

【资料变更】

请确认是否涉及资料变更。如涉及，需要在PR中体现，并简要说明修改内容。如不涉及，需填写“不涉及”

不涉及。

原因说明：本次测试的目标接口 torch.distributed.tensor._utils.compute_local_shape_and_global_offset 属于带有 _utils 前缀的内部私有 API。上游 PyTorch 官方未对其开放公开的原生独立文档。且该 API 逻辑为纯 Python 的坐标数学推导，不涉及 NPU 底层内存操作或算子调用，无特殊硬件约束，故不涉及对外公开资料的编写与变更。

【接口变更】

请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及，需要详细说明接口以及对应的变更内容，同时需要在资料中体现。如不涉及，需填写“不涉及”

不涉及。本次提交仅为纯测试层面的用例补充。

【功能验证】

说明测试场景，测试方法。如果本次测试方式与常规单元测试不同，请详细说明您的测试步骤
新增/变更内容是否已新增/适配UT测试用例看护，并补充测试自验证截图

测试场景与方法：在分布式环境下（世界大小为 8），测试传入 Partial 策略时的尺寸和偏移量推导是否正确。
执行步骤：

python test_utils.py -v -k test_compute_local_shape_and_global_offset_partial

测试结果：8 个通信进程均正常计算完毕并匹配正确断言，结果为 OK。

910B3的4卡测试自验证截图，无A5多卡环境，无法验证A5支持情况，不再进行A5多卡验证测试：

修改bug之后的函数测试：
执行步骤：

python test_utils.py -v -k test_compute_local_shape_2d_uneven

验证截图：

【CheckList】

PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]

• 代码注释完备，正确记录错误日志
• 代码实现进行了返回值、空指针等校验
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
• PR持续集成流水线（CI）执行通过，代码检查无异常