import torch
from torch_xla.stablehlo import exported_program_to_stablehlo
from torch_xla.experimental.mark_pattern_utils import StableHLOCompositeBuilder

class SubModule(torch.nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x, y):
        builder = StableHLOCompositeBuilder("abc.SubModule")
        x, y = builder.mark_inputs(x, y)
        out = x + y
        out = builder.mark_outputs(out)
        return out


class Model(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.submodule = SubModule()

    def forward(self, x, y):
        builder = StableHLOCompositeBuilder("abc.Model")
        x, y = builder.mark_inputs(x, y)
        a = x + y
        b = x - y
        c = self.submodule(a, b)
        a, b, c = builder.mark_outputs(a, b, c)
        return a + b + c 

sample_input = (torch.randn(1, 1, 32, 32), torch.randn(1, 1, 32, 32))
exported = torch.export.export(Model(), sample_input)
stablehlo_program = exported_program_to_stablehlo(exported)

module @IrToHlo.32 attributes {mhlo.cross_program_prefetches = [], mhlo.input_output_alias = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} {
  func.func @main(%arg0: tensor<1x1x32x32xf32>, %arg1: tensor<1x1x32x32xf32>) -> tensor<1x1x32x32xf32> {
    %0:3 = stablehlo.composite "abc.Model" %arg1, %arg0 {decomposition = @abc.Model.impl} : (tensor<1x1x32x32xf32>, tensor<1x1x32x32xf32>) -> (tensor<1x1x32x32xf32>, tensor<1x1x32x32xf32>, tensor<1x1x32x32xf32>)
    %1 = stablehlo.add %0#0, %0#1 : tensor<1x1x32x32xf32>
    %2 = stablehlo.add %1, %0#2 : tensor<1x1x32x32xf32>
    return %2 : tensor<1x1x32x32xf32>
  }
  func.func private @abc.SubModule.impl(%arg0: tensor<1x1x32x32xf32>, %arg1: tensor<1x1x32x32xf32>) -> tensor<1x1x32x32xf32> {
    %0 = stablehlo.add %arg0, %arg1 : tensor<1x1x32x32xf32>
    return %0 : tensor<1x1x32x32xf32>
  }
  func.func private @abc.Model.impl(%arg0: tensor<1x1x32x32xf32>, %arg1: tensor<1x1x32x32xf32>) -> (tensor<1x1x32x32xf32>, tensor<1x1x32x32xf32>, tensor<1x1x32x32xf32>) {
    %0 = stablehlo.add %arg0, %arg1 : tensor<1x1x32x32xf32>
    %1 = stablehlo.subtract %arg0, %arg1 : tensor<1x1x32x32xf32>
    %2 = stablehlo.composite "abc.SubModule" %0, %1 {decomposition = @abc.SubModule.impl} : (tensor<1x1x32x32xf32>, tensor<1x1x32x32xf32>) -> tensor<1x1x32x32xf32>
    return %0, %1, %2 : tensor<1x1x32x32xf32>, tensor<1x1x32x32xf32>, tensor<1x1x32x32xf32>
  }
}