experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Include new op EnqueueTPUEmbeddingRaggedTensorBatch in the list of TPUEmbedding enqueue ops don't get auto control edges created so that they can be used in TF2/eager mode in tf.functions.

Browse Source 
PiperOrigin-RevId: 310168476
Change-Id: I0149f613c8f407a774531a3c0219481b1d347db1
This commit is contained in:
Bruce Fontaine 2020-05-06 09:37:05 -07:00 committed by TensorFlower Gardener
parent dd00d506b3
commit ca946faa21
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/core/grappler/optimizers/function_optimizer.cc
View File

@ -849,7 +849,8 @@ const bool IsExemptFromSideEffectsExecutionValidation(const string& op) {
// TPUEmbedding EnqueueOps are stateful but this is only between ops with
// the same device_ordinal on the same host.
"EnqueueTPUEmbeddingSparseBatch", "EnqueueTPUEmbeddingIntegerBatch",
"EnqueueTPUEmbeddingSparseTensorBatch"});
"EnqueueTPUEmbeddingSparseTensorBatch",
"EnqueueTPUEmbeddingRaggedTensorBatch"});
return exemption->contains(op);
}

3
tensorflow/python/framework/auto_control_deps.py
View File

@ -99,7 +99,8 @@ LEGACY_RANDOM_OPS = [
_ORDER_INSENSITIVE_STATEFUL_OPS = [
"CudnnRNNV2", "CudnnRNNV3", "CudnnRNNBackpropV2", "CudnnRNNBackpropV3",
"EnqueueTPUEmbeddingSparseBatch", "EnqueueTPUEmbeddingIntegerBatch",
"EnqueueTPUEmbeddingSparseTensorBatch"
"EnqueueTPUEmbeddingSparseTensorBatch",
"EnqueueTPUEmbeddingRaggedTensorBatch"
]
# LINT.ThenChange(//tensorflow/core/grappler/optimizers/function_optimizer.cc)