Keras EffectiveTransformer layer.

Summary:
IPU implementation of https://github.com/bytedance/effective_transformer

TF2.4 Only

Test Plan:
Inference and training tests.
Unpadding test.
Arg tests.

Reviewers: #tensorflow, simonl, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, samuelm, nicolasc, georgep

Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, samuelm, nicolasc, georgep

Subscribers: jamiep, markf, mrinal, georgep

Maniphest Tasks: T35014

Differential Revision: https://phabricator.sourcevertex.net/D43654

Author: JackHunt

tensorflow/compiler/plugin/poplar/docs/api.rst

@@ -167,7 +167,7 @@ Keras layers
 
 
 .. automodule:: tensorflow.python.ipu.keras.layers
-  :members: Dropout, Embedding, GroupNormalization, InstanceNormalization, LayerNormalization, PopnnGRU, PopnnLSTM, SerialDense, CTCInferenceLayer, CTCPredictionsLayer, RecomputationCheckpoint
+  :members: Dropout, EffectiveTransformer, Embedding, GroupNormalization, InstanceNormalization, LayerNormalization, PopnnGRU, PopnnLSTM, SerialDense, CTCInferenceLayer, CTCPredictionsLayer, RecomputationCheckpoint
   :imported-members: Dropout, Embedding, GroupNormalization, InstanceNormalization, LayerNormalization, PopnnGRU, PopnnLSTM, SerialDense, CTCInferenceLayer, CTCPredictionsLayer, RecomputationCheckpoint
 
 Keras losses

tensorflow/python/ipu/BUILD

@@ -126,6 +126,7 @@ py_library(
         "keras/layers/assume_equal_across_replicas.py",
         "keras/layers/ctc.py",
         "keras/layers/dropout.py",
+        "keras/layers/effective_transformer.py",
         "keras/layers/embedding_lookup.py",
         "keras/layers/ipu_layer.py",
         "keras/layers/math.py",
@@ -538,6 +539,18 @@ tf_py_test(
     ],
 )
 
+tf_py_test(
+    name = "keras_effective_transformer_test",
+    size = "large",
+    srcs = ["tests/keras/keras_effective_transformer_test.py"],
+    shard_count = 12,
+    deps = [
+        "//tensorflow/python:framework_test_lib",
+        "//tensorflow/python/ipu:ipu_lib",
+        "//third_party/py/numpy",
+    ],
+)
+
 tf_py_test(
     name = "keras_mnist_test",
     size = "large",
@@ -1963,6 +1976,7 @@ test_suite(
         "keras_asynchronous_callbacks_test",
         "keras_ctc_test",
         "keras_dropout_test",
+        "keras_effective_transformer_test",
         "keras_embedding_lookup_test",
         "keras_execution_parameters_test",
         "keras_functional_model_test",

tensorflow/python/ipu/keras/layers/__init__.py

@@ -21,6 +21,7 @@
 from tensorflow.python.ipu.keras.layers.ctc import CTCInferenceLayer
 from tensorflow.python.ipu.keras.layers.ctc import CTCPredictionsLayer
 from tensorflow.python.ipu.keras.layers.dropout import Dropout
+from tensorflow.python.ipu.keras.layers.effective_transformer import EffectiveTransformer
 from tensorflow.python.ipu.keras.layers.embedding_lookup import Embedding
 from tensorflow.python.ipu.keras.layers.normalization import GroupNorm, InstanceNorm, LayerNorm
 from tensorflow.python.ipu.keras.layers.normalization import GroupNormalization, InstanceNormalization, LayerNormalization