%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>

  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

            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>(