Stablehlo Assertion `!impl->wasOpReplaced(op) && "attempting to modify a replaced/erased op"' failed.

[uazizTT]

Example graph:

%c = stablehlo.constant dense<1> : tensor<19xi64>
%0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<19xi64>) -> tensor<19xi64>
%1 = stablehlo.add %0, %c : tensor<19xi64>

Error dump:

ttmlir-opt: /localdev/aknezevic/tt-mlir/env/build/llvm-project-prefix/src/llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp:1589: virtual void mlir::ConversionPatternRewriter::startOpModification(mlir::Operation *): Assertion `!impl->wasOpReplaced(op) && "attempting to modify a replaced/erased op"' failed.

[mmanzoorTT]

The following PyTorch example generates the same error

def test_broadcast_add():
  class Basic(nn.Module):
    def __init__(self):
      super().__init__()

    def forward(self, x, y):
      return  x + y

  input1 = torch.tensor([1], dtype=torch.float32)
  input2 = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.float32)

  tt_mod = torch.compile(Basic(), backend=backend)
  out = tt_mod(input1, input2)
  print(f"output: {out}")

Generated StableHLO graph

module {
  func.func @main(%arg0: tensor<8xf32>, %arg1: tensor<1xf32>) -> tensor<8xf32> {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<8xf32>) -> tensor<8xf32>
    %1 = stablehlo.broadcast_in_dim %arg1, dims = [0] : (tensor<1xf32>) -> tensor<8xf32>
    %2 = stablehlo.add %0, %1 : tensor<8xf32>
    return %2 : tensor<8xf32>
  }
}

PyTorch is broadcasting both operands, although there is no need to broadcast %arg0 (output shape is same as input).