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>;
|
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", []> {
|
def TF__RecvTPUEmbeddingActivationsOp : TF_Op<"_RecvTPUEmbeddingActivations", []> {
|
||||||
let summary = "An op that receives embeddng activations on the TPU.";
|
let summary = "An op that receives embeddng activations on the TPU.";
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue