func.getLoc(), op.getOperandTypes(), /*dynamic_key=*/dynamic_key,
          op.getSendKeyAttr(),
          /*device_ordinal=*/rewriter.getI64IntegerAttr(0),
          rewriter.getStringAttr("TPU"));
      for (auto result :
           llvm::zip(cloned_func.getArguments(), recv_at_host->getResults())) {
        std::get<0>(result).replaceAllUsesWith(std::get<1>(result));
      }

  const TorchAvgPoolData pool = GetTorchAvgPoolData(op);

  if (pool.ph == 0 && pool.pw == 0) {
    return builder.getStringAttr("VALID");
  }
  if (pool.h_out == pool.h_in && pool.w_out == pool.w_in) {
    return builder.getStringAttr("SAME");
  }
  return builder.getStringAttr("CUSTOM");
}

}  // namespace odml

    TensorType tensor_type = mlir::cast<TensorType>(DropRefType(shaped_type));
    StringAttr device_attr =
        variable_v2_op->getAttrOfType<StringAttr>("device");
    if (!device_attr) device_attr = builder.getStringAttr("");
    StringRef variable_name =
        GetNodeNameFromClassAttrOrSharedNameAttr(variable_v2_op);
    if (variable_name.empty()) {
      return signalPassFailure();
    }

// satisfied by SAME or VALID tensorflow padding.
def HasCustomPadding:
  Constraint<CPred<"GetAvgPoolOpPadAttr($_builder, (*$0.begin()).getDefiningOp<mhlo::CompositeOp>()) == $_builder.getStringAttr(\"CUSTOM\")">>;

// Returns true if the provided padding in the composite op can be satisfied 
// by SAME or VALID tensorflow padding.
def HasSameOrValidPadding: Constraint<Neg<HasCustomPadding.predicate>>;

                                           device.getValue().str(), "_",
                                           inference_type.getValue().str());
  added_func_op.setName(builder.getStringAttr(function_name));
  added_call_op.setCallee(builder.getStringAttr(function_name));
}

// Raises partitioned sequential `Operations` from a block to a new function
// definition. `Operations` are partitioned into classes from the cartesian

                            live_ins);

  func::FuncOp outlined_func =
      BuildFunction(live_ins.getArrayRef(), cluster_op, symbol_table, builder);
  cluster_op->setAttr(
      builder->getStringAttr(kFuncAttr),
      mlir::SymbolRefAttr::get(builder->getContext(), outlined_func.getName()));

  builder->setInsertionPoint(cluster_op);
  auto cluster_func_op = builder->create<tf_device::ClusterFuncOp>(

        module_op.getLoc(), new_types, new_args,
        SymbolRefAttr::get(context, func_op.getSymName()),
        /*config=*/builder.getStringAttr(""),
        /*config_proto=*/builder.getStringAttr(""),
        /*executor_type=*/builder.getStringAttr(""));
    call_op_returns.append(call_op.getResults().begin(),
                           call_op.getResults().end());
    SetFunctionPrivate(func_op);
  }

        mlir::StringAttr num_layers_str = builder.getStringAttr("num_layers");
        NamedAttribute num_layers_attr(
            num_layers_str,
            IntegerAttr::get(IntegerType::get(fn.getContext(), /*width=*/32),
                             num_layers));

        // current_layer_index Composite Attribute.
        mlir::StringAttr current_layer_str =
            builder.getStringAttr("layer_index");

      location, output_types, args,
      FlatSymbolRefAttr::get(builder.getStringAttr(func_name)),
      /*config=*/"", /*config_proto=*/"", /*executor_type=*/"");

  // Set the attribute to annotate this function call op as a quantizable spot.
  call_op->setAttr(
      kQuantTraitAttrName,
      builder.getStringAttr(StringRef(
          std::string(QuantTraitValues[QuantizationTrait::FullyQuantizable]))));

def CustomOptions : NativeCodeCall<
  "CustomOptions($_builder.getContext(), $0.getDefiningOp()).value()">;

class CreateStringAttr<string values> : NativeCodeCall<
  "$_builder.getStringAttr("# values #")">;

def LegalizeTensorListReserve : Pat<(TF_TensorListReserveOp:$tf_op $shape, $num_elements),
    (TFL_CustomOp (Size2InputRange $shape, $num_elements),