1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
|
# mypy: allow-untyped-decorators
# Copyright (c) Meta Platforms, Inc. and affiliates
# implement matrix related ops for distributed tensor
import torch
from torch.distributed.tensor._dtensor_spec import DTensorSpec
from torch.distributed.tensor._op_schema import (
OpSchema,
OpStrategy,
PlacementStrategy,
StrategyType,
)
from torch.distributed.tensor._ops.utils import register_op_strategy
from torch.distributed.tensor.device_mesh import DeviceMesh
from torch.distributed.tensor.placement_types import Replicate
aten = torch.ops.aten
@register_op_strategy(aten.slice_backward.default)
def slice_backward_rules(mesh: DeviceMesh, op_schema: OpSchema) -> StrategyType:
"""
slice_backward is a new_zeros + slice_scatter, we only allow replication
on the input/output for now since new_zeros would produce replication
"""
replicate_spec = DTensorSpec(mesh, tuple([Replicate()] * mesh.ndim))
return OpStrategy([PlacementStrategy(replicate_spec)])
|
