Create _HostComputeMlir for use in TF MLIR.
PiperOrigin-RevId: 311792286 Change-Id: I6ec57f9b23c17dd52e756ead4ddfad58ecdb2f76
This commit is contained in:
parent
dc1c299833
commit
31583920dc
|
@ -10586,6 +10586,27 @@ def TF_ZerosLikeOp : TF_Op<"ZerosLike", [NoSideEffect, SameOperandsAndResultType
|
|||
TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
|
||||
}
|
||||
|
||||
def TF__HostComputeMlirOp : TF_Op<"_HostComputeMlir", []> {
|
||||
let summary = "A host-side computation called from a TPU device.";
|
||||
|
||||
let description = [{
|
||||
}];
|
||||
|
||||
let arguments = (ins
|
||||
Variadic<TF_Tensor>:$inputs,
|
||||
|
||||
StrAttr:$key,
|
||||
DefaultValuedAttr<I64Attr, "0">:$tpu_core
|
||||
);
|
||||
|
||||
let results = (outs
|
||||
Variadic<TF_Tensor>:$outputs
|
||||
);
|
||||
|
||||
TF_DerivedOperandTypeListAttr Tinputs = TF_DerivedOperandTypeListAttr<0>;
|
||||
TF_DerivedResultTypeListAttr Toutputs = TF_DerivedResultTypeListAttr<0>;
|
||||
}
|
||||
|
||||
def TF__RecvTPUEmbeddingActivationsOp : TF_Op<"_RecvTPUEmbeddingActivations", []> {
|
||||
let summary = "An op that receives embeddng activations on the TPU.";
|
||||
|
||||
|
|
Loading…
Reference in New Issue