(
value: Expr, device_mesh: py_str | DeviceMesh, placement: py_str | Placement
)
| 151 | |
| 152 | |
| 153 | def annotate_sharding( |
| 154 | value: Expr, device_mesh: py_str | DeviceMesh, placement: py_str | Placement |
| 155 | ) -> Expr: |
| 156 | if isinstance(device_mesh, py_str): |
| 157 | device_mesh = _lookup_device_mesh(device_mesh) |
| 158 | if isinstance(placement, py_str): |
| 159 | placement = Placement.from_text(placement) |
| 160 | return _annotate_sharding(value, device_mesh, placement) |
| 161 | |
| 162 | |
| 163 | def redistribute( |
nothing calls this directly
no test coverage detected
searching dependent graphs…