}

TEST_F(LiftAsFunctionCallTest, EinsumSupportedForXlaDotV2Succeeds) {
  StringAttr einsum_supported_by_xla_dot_v2_attr =
      builder_.getStringAttr("ijk,ikm->ijm");
  StringAttr einsum_one_operand = builder_.getStringAttr("ijk->ikj");
  StringAttr einsum_ellipsis = builder_.getStringAttr("...gse->...gs");
  EXPECT_TRUE(IsEinsumSupportedByXlaDotV2(einsum_supported_by_xla_dot_v2_attr));

  for (auto arg_attr : arg_attrs) {
    if (arg_attr.getName() == "tf_saved_model.bound_input") continue;
    new_arg_attrs.emplace_back(arg_attr);
  }

  const NamedAttribute index_path_attr(
      builder.getStringAttr(kTfSavedModelIndexPathAttr),
      builder.getStrArrayAttr({index_path}));

  new_arg_attrs.emplace_back(index_path_attr);
  return new_arg_attrs;
}

      DenseIntElementsAttr::get(
          RankedTensorType::get({1, num_replicas}, builder.getIntegerType(32)),
          group_assignment_val));

  StringAttr reduce_op = builder.getStringAttr("Add");
  StringAttr mode = builder.getStringAttr("CrossReplica");
  return builder.create<XlaAllReduceOp>(block_arg.getLoc(), block_arg.getType(),
                                        block_arg, group_assignment, reduce_op,

  OpBuilder builder(&ctx);
  // Set the insertion point right before the return op.
  builder.setInsertionPoint(&region.back().back());

  StringAttr output_file_path_attr = builder.getStringAttr(output_file_path);
  ArrayAttr ids_attr = builder.getStrArrayAttr(ids);
  ArrayAttr calibration_methods_attr =
      builder.getI32ArrayAttr(calibration_methods);
  builder.create<TF::CalibrationStatisticsSaverOp>(

  auto caller = builder.create<TF::StatefulPartitionedCallOp>(
      loc, result_types, operands, symbol,
      /*config=*/builder.getStringAttr(""),
      /*config_proto=*/builder.getStringAttr(""),
      /*executor_type=*/builder.getStringAttr(""));
  caller.setFAttr(symbol);

  // Set an attribute that our inliner will look for when choosing which
  // TF::StatefulPartitionedCallOps to inline.

      mean_op.getLoc(), avg_pool_output_type, input,
      rewriter.getI32IntegerAttr(height), rewriter.getI32IntegerAttr(width),
      rewriter.getStringAttr("VALID"), rewriter.getI32IntegerAttr(1),
      rewriter.getI32IntegerAttr(1), rewriter.getStringAttr("NONE"));

  auto value_to_replace = avg_pool.getResult();

  // If it's not keep dim, we need to insert a reshape after the average
  // pool.

              /*lhs=*/input,
              /*rhs=*/filter, window_strides, padding, lhs_dilation,
              rhs_dilation, feature_group_count,
              builder.getStringAttr(dnums.SerializeAsString()),
              /*precision_config=*/builder.getStringAttr(precision_config_str))
          .getOutput();

  // Dynamic-range quantization wil always fall into this case.
  if (input_zp_value == 0) return xla_conv_output;

  func->setAttr(kTfSavedModelExportedNamesAttr,
                builder.getStrArrayAttr({kSessionInitFuncName}));
  func->setAttr(kTfSavedModelInitializerTypeAttr,
                builder.getStringAttr(kTfSavedModelInitializerRestoreType));
  func.setVisibility(mlir::func::FuncOp::Visibility::Public);
  auto func_builder = OpBuilder::atBlockBegin(func.addEntryBlock());
  func_builder.create<mlir::func::ReturnOp>(func.getLoc());

class CreateI64ArrayAttr<string values> : NativeCodeCall<
  "$_builder.getI64ArrayAttr("# values #")">;

// Creates a string attribute with given values.
class CreateStringAttr<string values> : NativeCodeCall<
  "$_builder.getStringAttr(\""# values #"\")">;

// Creates a new F32 type with the same shape as the given value.
def CloneTypeWithF32ElementType : NativeCodeCall<
  "CloneTypeWithNewElementType($0.getType(), $_builder.getF32Type())">;

    return failure();
  }

  new_device_launch_op = builder.create<tf_device::LaunchOp>(
      tpu_copy_with_dynamic_shape_op->getLoc(),
      builder.getStringAttr(device_str),
      /*result_types=*/tpu_copy_with_dynamic_shape_op->getResultTypes());

  new_device_launch_op.getBody().push_back(new Block);
  builder.setInsertionPointToEnd(&new_device_launch_op.GetBody());