// Add the second argument to the first argument, which is expected to be an
// argument list.
// bias(einsum(inputs), bias) --> einsum_with_bias(AppendToVector(inputs, bias))
// Since inputs is a vector in case of einsum, we cannot use ArgumentList here.
def AppendToVector : NativeCodeCall<"AppendToVector($0, $1)">;

// The list of arguments of the composite function.
def ArgumentList : NativeCodeCall<"llvm::SmallVector<Value>{$0...}">;

def LiftEinsumWithBias : Pat<
  (TF_BiasAddOp:$res
    (TF_EinsumOp $input, $equation),
    $bias, IsDataFormatNHWC:$bias_data_format),
  (LiftAsTFPartitionedCall<"composite_einsum_with_bias_fn">
    (AppendToVector (ArgumentList $input), $bias),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"equation"> $equation))),
  [(IsNotInLiftedFunc $res),
   (IsEinsumSupportedByXlaDotV2 $equation)],

// Add the second argument to the first argument, which is expected to be an
// argument list.
// Used to attach bias to einsum argument list.
SmallVector<Value> AppendToVector(ArrayRef<Value> arguments, Value append);

// Checks if the `Method` attatched to the given `tf.XlaCallModule` op has
// `WeightOnlyPtq`.
bool HasWeightOnlyPtqMethod(TF::XlaCallModuleOp xla_call_module_op);

  SmallVector<NamedAttribute> attributes;
  return LiftAsFunctionCall(builder, location, call_op_type, func_name,
                            arguments, results, attributes);
}

SmallVector<Value> AppendToVector(const ArrayRef<Value> arguments,
                                  Value append) {
  SmallVector<Value> ret(arguments);
  ret.push_back(append);
  return ret;
}