from torch_xla.stablehlo import allow_custom_op_in_graph

m = Library("my_custom_library", "DEF")
m.define("custom_op(Tensor input) -> Tensor")

class M(torch.nn.Module):

  def forward(self, x):
    x = torch.sin(x)
    x = torch.ops.my_custom_library.custom_op(x)
    x = torch.cos(x)
    return x

options = StableHLOExportOptions()
options.custom_ops_allowed_in_graph.add("my_custom_library")
ep = torch.export.export(M(), (torch.randn(3, 3),))
shlo_module = exported_program_to_stablehlo(ep, options)

module @IrToHlo.6 attributes {mhlo.cross_program_prefetches = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} {
  func.func @main(%arg0: tensor<3x3xf32>) -> tensor<3x3xf32> {
    %0 = stablehlo.sine %arg0 : tensor<3x3xf32>
    %1 = stablehlo.custom_call @my_custom_library.custom_op.default(%0) {api_version = 0 : i32} : (tensor<3x3xf32>) -> tensor<3x3xf32>
    %2 = stablehlo.cosine %1 : tensor<3x3xf32>
    return %2 : tensor<3x3xf32>
  }
}

m.define("custom_op3(Tensor input) -> Tensor")

@impl(m, "custom_op3", "XLA")
def custom_op3_xla(x):
  res = stablehlo_custom_call((x,), "custom_op3", [x.shape[1:]],
                                  [torch.int8], True, "backend_config", 1)
  return res

class M(torch.nn.Module):

  def forward(self, x):
    x = torch.sin(x)
    x = torch.ops.my_custom_library.custom_op3(x)
    x = x + 1
    return x

ep = torch.export.export(M(), (torch.randn(3, 3),))
shlo_module = exported_program_to_stablehlo(ep)
shlo_text = shlo_module.get_stablehlo_text()

module @IrToHlo.10 attributes {mhlo.cross_program_prefetches = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} {
  func.func @main(%arg0: tensor<3x3xf32>) -> tensor<3xi8> {
    %c = stablehlo.constant dense<1> : tensor<3xi8>
    %0 = stablehlo.sine %arg0 : tensor<3x3xf32>
    %1 = stablehlo.custom_call @custom_op3(%0) {backend_config = "backend_config", has_side_effect = true} : (tensor<3x3xf32>) -> tensor<3xi8>
    %2 = stablehlo.add %1, %c : tensor<3xi8>
    return %2 : tensor<3xi8>
  }
}