rewriter.replaceOp(batch_norm, op->getResults());

    // Depending on the case, we may fuse the add, the relu, or both.
    if (!add_op || add_op.getZ().hasOneUse()) {
      // We fuse the Relu only if the add has a single use, otherwise we only
      // fuse the add itself.
      op->setAttr("activation_mode", rewriter.getStringAttr("Relu"));
      rewriter.replaceOp(relu_op, op->getResult(0));
    }
    if (add_op) {

  %cpu = corert.get_op_handler %ch "cpu"
  %0 = corert.executeop(%cpu) "tf.Relu"(%arg0) { T = f32 } : 1
  %arg1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %arg1_th {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
  %1 = tfrt_fallback_async.executeop key(0) cost(100) device("/CPU:0") "tf.Relu"(%arg1) { T = f32 } : 1
  tfrt.return

// CHECK: %[[cast:.*]] = "tf.Cast"(%[[sub]]) <{Truncate = false}> : (tensor<1x3x2x2xi32>) -> tensor<1x3x2x2xf32>
// CHECK: %[[dequant1:.*]] = "tf.Mul"(%[[cast]]
// CHECK: %[[relu:.*]] = "tf.Relu"(%[[dequant1]]
// CHECK: %[[clamped:.*]] = "tf.Minimum"(%[[relu]]

// CHECK: %[[rescale1:.*]] = "tf.Mul"(%[[cast]]
// CHECK: %[[add2:.*]] = "tf.AddV2"(%[[rescale1]]
// CHECK: %[[maximum2:.*]] = "tf.Maximum"(%[[add2]]

  %1 = "tfl.mul"(%0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
  func.return %1 : tensor<10x10x10xf32>
}

// -----

// CHECK: tac.cost = 0x4B673001

// CHECK-LABEL: @main

  op: "custom_embedding_matmul"
}
library {
  function {
    signature {
      name: "custom_relu"
    }
    attr {
      key: "_implements"
      value {
        func {
          name: "tensorflow.relu"
        }
      }
    }
  }
  function {
    signature {
      name: "custom_embedding_matmul"
    }
    attr {
      key: "_implements"
      value {
        func {

  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
  // CHECK: tac.cost = 1.000000e+03
  %1 = "tfl.mul"(%0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
  func.return %1 : tensor<10x10x10xf32>
}

// -----

  %14 = "tf.AddV2"(%10, %12) : (tensor<?x256x56x56xf32>, tensor<?x256x56x56xf32>) -> tensor<?x256x56x56xf32>
  %15 = "tf.Relu"(%14) : (tensor<?x256x56x56xf32>) -> tensor<?x256x56x56xf32>

  // CHECK: %[[ADD:[0-9]*]] = "tf.AddV2"(%[[BATCH_NORM1]], %[[BATCH_NORM2]])
  // CHECK: %[[RELU:[0-9]*]] = "tf.Relu"(%[[ADD]])

  // Reduce spatial dimensions
  %16 = "tf.Mean"(%15, %1) : (tensor<?x256x56x56xf32>, tensor<2xi32>) -> tensor<?x256xf32>

  // CHECK: device = "cpu"
  %2 = "tf.Relu"(%1) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "cpu"} : (tensor<3x3xf32>) -> tensor<3x3xf32>
  // CHECK: device = "gpu"
  %3 = "tf.Relu"(%2) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"]} : (tensor<3x3xf32>) -> tensor<3x3xf32>
  func.return %3 : tensor<3x3xf32>

    value {
      s: "NHWC"
    }
  }
}
node {
  name: "conv_net_2d_1/Relu"
  op: "Relu"
  input: "conv_net_2d_1/conv_2d_0/BiasAdd"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
}
node {
  name: "output_0"
  op: "Identity"
  input: "conv_net_2d_1/Relu"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
}