Currently, distconv is not compatible with torch.compile. But it seems pytorch tensor subclasses (such as DTensor) do support this. We should find a way to add this capability to distconv.
I think it is through __tensor_flatten__ and __tensor_unflatten__ methods.
Currently,
distconvis not compatible withtorch.compile. But it seems pytorch tensor subclasses (such asDTensor) do support this. We should find a way to add this capability todistconv.I think it is through
__tensor_flatten__and__tensor_unflatten__methods.