%1 = "tfl.shape"(%conv_s) : (tensor<?x26x26x26x16xf32>) -> tensor<5xi64>
  %2 = "tfl.broadcast_args"(%1, %cst) : (tensor<5xi64>, tensor<1xi64>) -> tensor<5xi64>
  %broad1 = "tfl.broadcast_to"(%conv_s, %2) : (tensor<?x26x26x26x16xf32>, tensor<5xi64>) -> tensor<?x26x26x26x16xf32>

  // CHECK: %[[BROADCAST_ARGS_1:.*]] = "tfl.broadcast_args"(%[[BROADCAST_ARGS_0]], %[[SHAPE_2]]) : (tensor<8xi64>, tensor<8xi64>) -> tensor<8xi64>
  // CHECK: %[[BROADCAST_TO_2:.*]] = "tfl.broadcast_to"(%[[LESS]], %[[BROADCAST_ARGS_1]]) : (tensor<8x7x6x5x?x3x2x1xi1>, tensor<8xi64>) -> tensor<8x7x6x5x?x3x2x1xi1>

  func.return %0 : tensor<?x?x?x?x?xf32>
}

// -----

func.func @broadcast_args(%arg0: tensor<5xi32>, %arg1: tensor<2xi32>) -> tensor<5xi32> {
  // CHECK: "tfl.broadcast_args"(%arg0, %arg1)
  %0 = "tfl.broadcast_args"(%arg0, %arg1) {} : (tensor<5xi32>, tensor<2xi32>) -> tensor<5xi32>
  func.return %0 : tensor<5xi32>
}

// -----

  HASHTABLE_FIND = 137,
  HASHTABLE_IMPORT = 138,
  HASHTABLE_SIZE = 139,
  REDUCE_ALL = 140,
  CONV_3D_TRANSPOSE = 141,
  VAR_HANDLE = 142,
  READ_VARIABLE = 143,
  ASSIGN_VARIABLE = 144,
  BROADCAST_ARGS = 145,
  RANDOM_STANDARD_NORMAL = 146,
  BUCKETIZE = 147,
  RANDOM_UNIFORM = 148,
  MULTINOMIAL = 149,
  GELU = 150,
  DYNAMIC_UPDATE_SLICE = 151,
  RELU_0_TO_1 = 152,
  UNSORTED_SEGMENT_PROD = 153,

                                             broadcast_multiplier);
    auto broadcast_rhs = broadcastToFeatureDim(
        bn_op.getLoc(), input_type, rhs, shape_value, feature_dim, rewriter);
    rewriter.replaceOpWithNewOp<mhlo::AddOp>(bn_op, lhs, broadcast_rhs);

    return success();
  }
};

class UnfuseMhloBatchNormPass

  let arguments = (ins
    TFL_ResourceTensor:$hash_table
  );

  let results = (outs
    TFL_I64Tensor:$out
  );
}

def TFL_BroadcastArgsOp : TFL_Op<"broadcast_args",[
    OperandsSameElementTypeConstraintBase<"BroadcastArgs op">,
    PredOpTrait<"input and output must have same element type",
      TFL_TCresVTEtIsSameAsOp<0, 0>>,
    TFL_OperandHasRank<0, 1>,
    TFL_OperandHasRank<1, 1>,

            mlir::cast<ShapedType>(rhs.getType()).getShape(),
            symbolic_broadcast_shape)) {
      return failure();
    }

    // Calculates the broadcast shape using BroadcastArgs op.
    Value lhs_shape = GetShape(lhs, op->getLoc(), rewriter);
    Value rhs_shape = GetShape(rhs, op->getLoc(), rewriter);
    auto broadcast_shape =
        rewriter
            .create<TF::BroadcastArgsOp>(

    "HASHTABLE_FIND",
    "HASHTABLE_IMPORT",
    "HASHTABLE_SIZE",
    "REDUCE_ALL",
    "CONV_3D_TRANSPOSE",
    "VAR_HANDLE",
    "READ_VARIABLE",
    "ASSIGN_VARIABLE",
    "BROADCAST_ARGS",
    "RANDOM_STANDARD_NORMAL",
    "BUCKETIZE",
    "RANDOM_UNIFORM",
    "MULTINOMIAL",
    "GELU",
    "DYNAMIC_UPDATE_SLICE",
    "RELU_0_TO_1",
    "UNSORTED_SEGMENT_PROD",

     {"ApproxTopK", "BroadcastTo", "ExpandDims", "Fill", "NoOp",
      "Range", "Rank", "Reshape", "Shape", "ShapeN", "Size", "Squeeze",
      "Transpose", "ZerosLike", "OnesLike", "BiasAdd" /*PW + Broadcast*/,
      "BroadcastArgs", "BroadcastGradientArgs", "OneHot", "Concat", "ConcatV2",
      "ConcatOffset", "Const", "MirrorPad", "MirrorPadGrad", "Pack", "Pad",
      "PadV2", "Reverse", "ReverseV2", "ReverseSequence", "Slice", "Split",

  // maximum rank.
  if (max_rank <= max_bcast_rank) {
    return true;
  }

  // Checks if all operands are broadcasted by BroadcastTo ops with the shape
  // is calculated from the same BroadcastArgs op. In such case, all operands
  // will have the same shape.
  Operation* broadcast_args_pivot = nullptr;
  for (unsigned index : indices) {
    Operation* parent_broadcast_args =