code: def test_amp_fuse_with_branch():
      class TestNet(nn.HybridBlock):
          def __init__(self, **kwargs):
              super(TestNet, self).__init__(**kwargs)
              self.fc1 = nn.Dense(16)
              self.fc2 = nn.Dense(16)
  
          def forward(self, x, *args):
              out = self.fc1(x)
              out1 = self.fc2(out)
              out1 = nn.Activation("relu")(out1)
              with nn.HybridBlock.OptConstraint.disable_amp():
                  out2 = mx.npx.softmax(out)
              return out1, out2
  
      net = TestNet()
      net.initialize()
  
      data_example = mx.np.ones((10,))
  
      #              |---> lp16_op_2
      # lp16_op_1 ---|
      #              |---> f32_amp_cast ---> f32_op
      #
      # `lp16_op_1` cannot fuse the `f32_amp_cast` node, since `lp16_op_2` already uses its lp16 output
  
      exp_data = mx.sym.Variable("data")
      exp_weight = [
          mx.symbol.Variable("weight{}".format(i)) for i in range(2)
      ]
      exp_bias = [
          mx.symbol.Variable("bias{}".format(i)) for i in range(2)
      ]
      amp_data = mx.sym.amp_cast(exp_data, dtype=AMP_DTYPE)
      lp16_op_1 = mx.sym.FullyConnected(
          amp_data, exp_weight[0], exp_bias[0], num_hidden=16
      )
      lp16_op_2 = mx.sym.FullyConnected(
          lp16_op_1, exp_weight[1], exp_bias[1], num_hidden=16
      )
      f32_amp_cast = mx.sym.amp_cast(lp16_op_1, dtype="float32")
      f32_op = mx.sym.softmax(f32_amp_cast)
      exp_sym = mx.sym.Group([lp16_op_2, f32_op])
      exp_sym = exp_sym.get_backend_symbol(SG_PASS_NAME)
      check_amp_fuse(net, [data_example], exp_sym)
  
file_path: mxnet/tests/python/dnnl/subgraphs/test_amp_subgraph.py
imports: import json
  import mxnet
  from mxnet import amp
  from mxnet.gluon import nn
  from mxnet.test_utils import assert_almost_equal
  from subgraph_common import SG_PASS_NAME, QUANTIZE_SG_PASS_NAME
  from test_matmul_subgraph import MultiHeadAttention
  import os
  import sys
  from pathlib import Path
  from amp.common import check_amp_net_stats
label_desc: Use image node_id 7 for calling a global function with example usage: test_amp_fuse_with_branch() and returns: out1,  out2