transformer.py 20 KB
Newer Older
Frederick Liu's avatar
Frederick Liu committed
1
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Hongkun Yu's avatar
Hongkun Yu committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Frederick Liu's avatar
Frederick Liu committed
14

Hongkun Yu's avatar
Hongkun Yu committed
15
"""Keras-based transformer block layer."""
16
# pylint: disable=g-classes-have-attributes
Hongkun Yu's avatar
Hongkun Yu committed
17

Chen Chen's avatar
Chen Chen committed
18
import gin
Hongkun Yu's avatar
Hongkun Yu committed
19
20
import tensorflow as tf

21
from official.nlp import keras_nlp
Hongkun Yu's avatar
Hongkun Yu committed
22
from official.nlp.modeling.layers import attention
23
from official.nlp.modeling.layers import multi_channel_attention
24
from official.nlp.modeling.layers.util import tf_function_if_eager
Hongkun Yu's avatar
Hongkun Yu committed
25
26
27


@tf.keras.utils.register_keras_serializable(package="Text")
28
class Transformer(keras_nlp.layers.TransformerEncoderBlock):
Hongkun Yu's avatar
Hongkun Yu committed
29
30
31
32
33
  """Transformer layer.

  This layer implements the Transformer from "Attention Is All You Need".
  (https://arxiv.org/abs/1706.03762).

34
  Args:
Hongkun Yu's avatar
Hongkun Yu committed
35
36
37
38
39
    num_attention_heads: Number of attention heads.
    intermediate_size: Size of the intermediate layer.
    intermediate_activation: Activation for the intermediate layer.
    dropout_rate: Dropout probability for the post-attention and output dropout.
    attention_dropout_rate: Dropout probability for within the attention layer.
40
41
    output_range: the sequence output range, [0, output_range) by slicing the
      target sequence. `None` means the target sequence is not sliced.
Hongkun Yu's avatar
Hongkun Yu committed
42
43
44
45
46
47
48
    kernel_initializer: Initializer for dense layer kernels.
    bias_initializer: Initializer for dense layer biases.
    kernel_regularizer: Regularizer for dense layer kernels.
    bias_regularizer: Regularizer for dense layer biases.
    activity_regularizer: Regularizer for dense layer activity.
    kernel_constraint: Constraint for dense layer kernels.
    bias_constraint: Constraint for dense layer kernels.
xinliupitt's avatar
xinliupitt committed
49
50
    use_bias: Whether to enable use_bias in attention layer. If set False,
      use_bias in attention layer is disabled.
xinliupitt's avatar
xinliupitt committed
51
    norm_first: Whether to normalize inputs to attention and intermediate dense
52
53
      layers. If set False, output of attention and intermediate dense layers is
      normalized.
xinliupitt's avatar
xinliupitt committed
54
    norm_epsilon: Epsilon value to initialize normalization layers.
55
    intermediate_dropout: Dropout probability for intermediate_dropout_layer.
xinliupitt's avatar
xinliupitt committed
56
57
    attention_initializer: Initializer for kernels of attention layers. If set
      `None`, attention layers use kernel_initializer as initializer for kernel.
Hongkun Yu's avatar
Hongkun Yu committed
58
59
60
61
62
63
64
65
  """

  def __init__(self,
               num_attention_heads,
               intermediate_size,
               intermediate_activation,
               dropout_rate=0.0,
               attention_dropout_rate=0.0,
66
               output_range=None,
Hongkun Yu's avatar
Hongkun Yu committed
67
68
69
70
71
72
73
               kernel_initializer="glorot_uniform",
               bias_initializer="zeros",
               kernel_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               kernel_constraint=None,
               bias_constraint=None,
xinliupitt's avatar
xinliupitt committed
74
75
76
               use_bias=True,
               norm_first=False,
               norm_epsilon=1e-12,
xinliupitt's avatar
xinliupitt committed
77
               intermediate_dropout=0.0,
xinliupitt's avatar
xinliupitt committed
78
               attention_initializer=None,
Hongkun Yu's avatar
Hongkun Yu committed
79
               **kwargs):
80
    super().__init__(
Zhenyu Tan's avatar
Zhenyu Tan committed
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
        num_attention_heads=num_attention_heads,
        inner_dim=intermediate_size,
        inner_activation=intermediate_activation,
        output_dropout=dropout_rate,
        attention_dropout=attention_dropout_rate,
        output_range=output_range,
        kernel_initializer=kernel_initializer,
        bias_initializer=bias_initializer,
        kernel_regularizer=kernel_regularizer,
        bias_regularizer=bias_regularizer,
        activity_regularizer=activity_regularizer,
        kernel_constraint=kernel_constraint,
        bias_constraint=bias_constraint,
        use_bias=use_bias,
        norm_first=norm_first,
        norm_epsilon=norm_epsilon,
        inner_dropout=intermediate_dropout,
        attention_initializer=attention_initializer,
        **kwargs)
100

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
101
102
103
104
105
106
107
108
109
  def get_config(self):
    return {
        "num_attention_heads":
            self._num_heads,
        "intermediate_size":
            self._inner_dim,
        "intermediate_activation":
            self._inner_activation,
        "dropout_rate":
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
110
            self._output_dropout_rate,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
        "attention_dropout_rate":
            self._attention_dropout_rate,
        "output_range":
            self._output_range,
        "kernel_initializer":
            tf.keras.initializers.serialize(self._kernel_initializer),
        "bias_initializer":
            tf.keras.initializers.serialize(self._bias_initializer),
        "kernel_regularizer":
            tf.keras.regularizers.serialize(self._kernel_regularizer),
        "bias_regularizer":
            tf.keras.regularizers.serialize(self._bias_regularizer),
        "activity_regularizer":
            tf.keras.regularizers.serialize(self._activity_regularizer),
        "kernel_constraint":
            tf.keras.constraints.serialize(self._kernel_constraint),
        "bias_constraint":
            tf.keras.constraints.serialize(self._bias_constraint),
        "use_bias":
            self._use_bias,
        "norm_first":
            self._norm_first,
        "norm_epsilon":
            self._norm_epsilon,
        "intermediate_dropout":
            self._inner_dropout,
        "attention_initializer":
            tf.keras.initializers.serialize(self._attention_initializer)
    }

141

Chen Chen's avatar
Chen Chen committed
142
143
@tf.keras.utils.register_keras_serializable(package="Text")
@gin.configurable
144
145
146
147
class CompiledTransformer(Transformer):

  @tf_function_if_eager(experimental_compile=True)
  def call(self, inputs):
148
    return super().call(inputs)
149
150
151


@tf.keras.utils.register_keras_serializable(package="Text")
Hongkun Yu's avatar
Hongkun Yu committed
152
class TransformerDecoderBlock(tf.keras.layers.Layer):
153
154
155
156
157
158
  """Single transformer layer for decoder.

  It has three sub-layers:
  (1) a multi-head self-attention mechanism.
  (2) a encoder-decoder attention.
  (3) a positionwise fully connected feed-forward network.
Hongkun Yu's avatar
Hongkun Yu committed
159

160
  Args:
Hongkun Yu's avatar
Hongkun Yu committed
161
162
163
164
165
166
167
168
169
170
171
172
173
174
    num_attention_heads: Number of attention heads.
    intermediate_size: Size of the intermediate layer.
    intermediate_activation: Activation for the intermediate layer.
    dropout_rate: Dropout probability for the post-attention and output dropout.
    attention_dropout_rate: Dropout probability for within the attention layer.
    multi_channel_cross_attention: Whether to use `MultiChannelAttention` for
      cross-attention between target sequences and source sequences.
    kernel_initializer: Initializer for dense layer kernels.
    bias_initializer: Initializer for dense layer biases.
    kernel_regularizer: Regularizer for dense layer kernels.
    bias_regularizer: Regularizer for dense layer biases.
    activity_regularizer: Regularizer for dense layer activity.
    kernel_constraint: Constraint for dense layer kernels.
    bias_constraint: Constraint for dense layer kernels.
xinliupitt's avatar
xinliupitt committed
175
176
    use_bias: Whether to enable use_bias in attention layer. If set False,
      use_bias in attention layer is disabled.
xinliupitt's avatar
xinliupitt committed
177
    norm_first: Whether to normalize inputs to attention and intermediate dense
178
179
      layers. If set False, output of attention and intermediate dense layers is
      normalized.
xinliupitt's avatar
xinliupitt committed
180
    norm_epsilon: Epsilon value to initialize normalization layers.
181
    intermediate_dropout: Dropout probability for intermediate_dropout_layer.
xinliupitt's avatar
xinliupitt committed
182
183
    attention_initializer: Initializer for kernels of attention layers. If set
      `None`, attention layers use kernel_initializer as initializer for kernel.
184
185
186
  """

  def __init__(self,
Hongkun Yu's avatar
Hongkun Yu committed
187
188
189
190
191
               num_attention_heads,
               intermediate_size,
               intermediate_activation,
               dropout_rate=0.0,
               attention_dropout_rate=0.0,
192
               multi_channel_cross_attention=False,
Hongkun Yu's avatar
Hongkun Yu committed
193
194
195
196
197
198
199
               kernel_initializer="glorot_uniform",
               bias_initializer="zeros",
               kernel_regularizer=None,
               bias_regularizer=None,
               activity_regularizer=None,
               kernel_constraint=None,
               bias_constraint=None,
xinliupitt's avatar
xinliupitt committed
200
201
202
               use_bias=True,
               norm_first=False,
               norm_epsilon=1e-12,
xinliupitt's avatar
xinliupitt committed
203
               intermediate_dropout=0.0,
xinliupitt's avatar
xinliupitt committed
204
               attention_initializer=None,
205
               **kwargs):
Hongkun Yu's avatar
Hongkun Yu committed
206
    super().__init__(**kwargs)
207
208
209
210
    self.num_attention_heads = num_attention_heads
    self.intermediate_size = intermediate_size
    self.intermediate_activation = tf.keras.activations.get(
        intermediate_activation)
Hongkun Yu's avatar
Hongkun Yu committed
211
212
    self.dropout_rate = dropout_rate
    self.attention_dropout_rate = attention_dropout_rate
213
    self.multi_channel_cross_attention = multi_channel_cross_attention
Hongkun Yu's avatar
Hongkun Yu committed
214
215
216
217
218
219
220
    self._kernel_initializer = tf.keras.initializers.get(kernel_initializer)
    self._bias_initializer = tf.keras.initializers.get(bias_initializer)
    self._kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
    self._bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
    self._activity_regularizer = tf.keras.regularizers.get(activity_regularizer)
    self._kernel_constraint = tf.keras.constraints.get(kernel_constraint)
    self._bias_constraint = tf.keras.constraints.get(bias_constraint)
xinliupitt's avatar
xinliupitt committed
221
222
223
    self._use_bias = use_bias
    self._norm_first = norm_first
    self._norm_epsilon = norm_epsilon
xinliupitt's avatar
xinliupitt committed
224
    self._intermediate_dropout = intermediate_dropout
xinliupitt's avatar
xinliupitt committed
225
    if attention_initializer:
xinliupitt's avatar
xinliupitt committed
226
227
      self._attention_initializer = tf.keras.initializers.get(
          attention_initializer)
xinliupitt's avatar
xinliupitt committed
228
229
    else:
      self._attention_initializer = self._kernel_initializer
230
231
232
233
234
    if self.multi_channel_cross_attention:
      self._cross_attention_cls = multi_channel_attention.MultiChannelAttention
    else:
      self._cross_attention_cls = attention.MultiHeadAttention

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
235
236
237
  def _maybe_build(self, inputs):
    super()._maybe_build(inputs[:1])

Hongkun Yu's avatar
Hongkun Yu committed
238
239
  def build(self, input_shape):
    target_tensor_shape = tf.TensorShape(input_shape[0])
240
    if len(target_tensor_shape.as_list()) != 3:
Hongkun Yu's avatar
Hongkun Yu committed
241
242
243
244
      raise ValueError("TransformerLayer expects a three-dimensional input of "
                       "shape [batch, sequence, width].")
    hidden_size = target_tensor_shape[2]
    if hidden_size % self.num_attention_heads != 0:
245
246
      raise ValueError(
          "The hidden size (%d) is not a multiple of the number of attention "
Hongkun Yu's avatar
Hongkun Yu committed
247
          "heads (%d)" % (hidden_size, self.num_attention_heads))
248
    self.attention_head_size = int(hidden_size) // self.num_attention_heads
249
    common_kwargs = dict(
250
        bias_initializer=self._bias_initializer,
Hongkun Yu's avatar
Hongkun Yu committed
251
252
253
254
        kernel_regularizer=self._kernel_regularizer,
        bias_regularizer=self._bias_regularizer,
        activity_regularizer=self._activity_regularizer,
        kernel_constraint=self._kernel_constraint,
255
256
257
258
        bias_constraint=self._bias_constraint)
    # Self attention.
    self.self_attention = attention.CachedAttention(
        num_heads=self.num_attention_heads,
259
        key_dim=self.attention_head_size,
260
        dropout=self.attention_dropout_rate,
xinliupitt's avatar
xinliupitt committed
261
        use_bias=self._use_bias,
xinliupitt's avatar
xinliupitt committed
262
        kernel_initializer=self._attention_initializer,
263
264
265
266
267
268
        name="self_attention",
        **common_kwargs)
    self.self_attention_output_dense = tf.keras.layers.experimental.EinsumDense(
        "abc,cd->abd",
        output_shape=(None, hidden_size),
        bias_axes="d",
xinliupitt's avatar
xinliupitt committed
269
        kernel_initializer=self._kernel_initializer,
270
271
        name="output",
        **common_kwargs)
272
    self.self_attention_dropout = tf.keras.layers.Dropout(
Hongkun Yu's avatar
Hongkun Yu committed
273
        rate=self.dropout_rate)
274
275
    self.self_attention_layer_norm = (
        tf.keras.layers.LayerNormalization(
xinliupitt's avatar
xinliupitt committed
276
            name="self_attention_layer_norm",
277
            axis=-1,
278
279
            epsilon=self._norm_epsilon,
            dtype="float32"))
280
281
282
    # Encoder-decoder attention.
    self.encdec_attention = self._cross_attention_cls(
        num_heads=self.num_attention_heads,
283
        key_dim=self.attention_head_size,
Hongkun Yu's avatar
Hongkun Yu committed
284
285
        dropout=self.attention_dropout_rate,
        output_shape=hidden_size,
xinliupitt's avatar
xinliupitt committed
286
        use_bias=self._use_bias,
xinliupitt's avatar
xinliupitt committed
287
        kernel_initializer=self._attention_initializer,
288
289
        name="attention/encdec",
        **common_kwargs)
290
291

    self.encdec_attention_dropout = tf.keras.layers.Dropout(
Hongkun Yu's avatar
Hongkun Yu committed
292
        rate=self.dropout_rate)
293
294
    self.encdec_attention_layer_norm = (
        tf.keras.layers.LayerNormalization(
xinliupitt's avatar
xinliupitt committed
295
            name="attention/encdec_output_layer_norm",
296
            axis=-1,
297
298
            epsilon=self._norm_epsilon,
            dtype="float32"))
299
300

    # Feed-forward projection.
301
302
303
304
    self.intermediate_dense = tf.keras.layers.experimental.EinsumDense(
        "abc,cd->abd",
        output_shape=(None, self.intermediate_size),
        bias_axes="d",
xinliupitt's avatar
xinliupitt committed
305
        kernel_initializer=self._kernel_initializer,
306
307
        name="intermediate",
        **common_kwargs)
308
309
    self.intermediate_activation_layer = tf.keras.layers.Activation(
        self.intermediate_activation)
310
311
    self._intermediate_dropout_layer = tf.keras.layers.Dropout(
        rate=self._intermediate_dropout)
312
313
314
315
    self.output_dense = tf.keras.layers.experimental.EinsumDense(
        "abc,cd->abd",
        output_shape=(None, hidden_size),
        bias_axes="d",
xinliupitt's avatar
xinliupitt committed
316
        kernel_initializer=self._kernel_initializer,
317
318
        name="output",
        **common_kwargs)
Hongkun Yu's avatar
Hongkun Yu committed
319
    self.output_dropout = tf.keras.layers.Dropout(rate=self.dropout_rate)
320
    self.output_layer_norm = tf.keras.layers.LayerNormalization(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
321
322
323
324
        name="output_layer_norm",
        axis=-1,
        epsilon=self._norm_epsilon,
        dtype="float32")
Hongkun Yu's avatar
Hongkun Yu committed
325
    super().build(input_shape)
326

xinliupitt's avatar
xinliupitt committed
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
  def get_config(self):
    config = {
        "num_attention_heads":
            self.num_attention_heads,
        "intermediate_size":
            self.intermediate_size,
        "intermediate_activation":
            self.intermediate_activation,
        "dropout_rate":
            self.dropout_rate,
        "attention_dropout_rate":
            self.attention_dropout_rate,
        "multi_channel_cross_attention":
            self.multi_channel_cross_attention,
        "kernel_initializer":
            tf.keras.initializers.serialize(self._kernel_initializer),
        "bias_initializer":
            tf.keras.initializers.serialize(self._bias_initializer),
        "kernel_regularizer":
            tf.keras.regularizers.serialize(self._kernel_regularizer),
        "bias_regularizer":
            tf.keras.regularizers.serialize(self._bias_regularizer),
        "activity_regularizer":
            tf.keras.regularizers.serialize(self._activity_regularizer),
        "kernel_constraint":
            tf.keras.constraints.serialize(self._kernel_constraint),
        "bias_constraint":
            tf.keras.constraints.serialize(self._bias_constraint),
        "use_bias":
            self._use_bias,
        "norm_first":
            self._norm_first,
        "norm_epsilon":
xinliupitt's avatar
xinliupitt committed
360
361
            self._norm_epsilon,
        "intermediate_dropout":
xinliupitt's avatar
xinliupitt committed
362
363
            self._intermediate_dropout,
        "attention_initializer":
xinliupitt's avatar
xinliupitt committed
364
            tf.keras.initializers.serialize(self._attention_initializer)
xinliupitt's avatar
xinliupitt committed
365
    }
Hongkun Yu's avatar
Hongkun Yu committed
366
    base_config = super().get_config()
xinliupitt's avatar
xinliupitt committed
367
368
    return dict(list(base_config.items()) + list(config.items()))

369
370
371
372
373
374
375
  def common_layers_with_encoder(self):
    """Gets layer objects that can make a Transformer encoder block."""
    return [
        self.self_attention, self.self_attention_layer_norm,
        self.intermediate_dense, self.output_dense, self.output_layer_norm
    ]

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
376
377
378
  def _parse_inputs(self, inputs, multi_channel_cross_attention):
    if multi_channel_cross_attention:
      if len(inputs) < 5:
379
        raise ValueError(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
380
            "TransformerDecoderBlock must have at least 5 inputs, when it uses "
381
            "multi_channel_cross_attention. But it got: %d" % len(inputs))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
      elif len(inputs) == 5:
        input_tensor, memory, attention_mask, self_attention_mask, context_attention_weights = inputs
        input_pos_embed = None
        memory_pos_embed = None
      elif len(inputs) == 6:
        input_tensor, memory, attention_mask, self_attention_mask, context_attention_weights, input_pos_embed = inputs
        memory_pos_embed = None
      else:
        input_tensor, memory, attention_mask, self_attention_mask, context_attention_weights, input_pos_embed, memory_pos_embed = inputs[:
                                                                                                                                         7]
    else:
      context_attention_weights = None
      if len(inputs) < 4:
        raise ValueError(
            "TransformerDecoderBlock must have at leaset 4 inputs, but it "
            "got: %d" % len(inputs))
      elif len(inputs) == 4:
        input_tensor, memory, attention_mask, self_attention_mask = inputs
        input_pos_embed = None
        memory_pos_embed = None
      elif len(inputs) == 5:
        input_tensor, memory, attention_mask, self_attention_mask, input_pos_embed = inputs
        memory_pos_embed = None
      else:
        input_tensor, memory, attention_mask, self_attention_mask, input_pos_embed, memory_pos_embed = inputs[:
                                                                                                              6]

    return input_tensor, memory, attention_mask, self_attention_mask, context_attention_weights, input_pos_embed, memory_pos_embed

  def call(self, inputs, cache=None, decode_loop_step=None):
    input_tensor, memory, attention_mask, self_attention_mask, context_attention_weights, input_pos_embed, memory_pos_embed = self._parse_inputs(
        inputs, self.multi_channel_cross_attention)

xinliupitt's avatar
xinliupitt committed
415
416
417
    source_tensor = input_tensor
    if self._norm_first:
      input_tensor = self.self_attention_layer_norm(input_tensor)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
418
419
420
421
422
423
    if input_pos_embed is None:
      self_attn_query = input_tensor
      self_attn_key = input_tensor
    else:
      self_attn_query = input_tensor + input_pos_embed
      self_attn_key = input_tensor + input_pos_embed
424
    self_attention_output, cache = self.self_attention(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
425
426
        query=self_attn_query,
        key=self_attn_key,
427
        value=input_tensor,
428
429
430
431
        attention_mask=self_attention_mask,
        cache=cache,
        decode_loop_step=decode_loop_step)
    self_attention_output = self.self_attention_dropout(self_attention_output)
xinliupitt's avatar
xinliupitt committed
432
433
434
435
436
437
438
439
440
    if self._norm_first:
      self_attention_output = source_tensor + self_attention_output
    else:
      self_attention_output = self.self_attention_layer_norm(
          input_tensor + self_attention_output)
    if self._norm_first:
      source_self_attention_output = self_attention_output
      self_attention_output = self.encdec_attention_layer_norm(
          self_attention_output)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
441
442
443
444
445
446
447
448
    if input_pos_embed is None:
      cross_attn_query = self_attention_output
    else:
      cross_attn_query = self_attention_output + input_pos_embed
    if memory_pos_embed is None:
      cross_attn_key = memory
    else:
      cross_attn_key = memory + memory_pos_embed
449
    cross_attn_inputs = dict(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
450
451
        query=cross_attn_query,
        key=cross_attn_key,
452
453
        value=memory,
        attention_mask=attention_mask)
454
455
    if self.multi_channel_cross_attention:
      # Accesses the 5-th input tensor for the doc-attention probabilities.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
456
      cross_attn_inputs["context_attention_weights"] = context_attention_weights
457
    attention_output = self.encdec_attention(**cross_attn_inputs)
458
    attention_output = self.encdec_attention_dropout(attention_output)
xinliupitt's avatar
xinliupitt committed
459
460
461
462
    if self._norm_first:
      attention_output = source_self_attention_output + attention_output
    else:
      attention_output = self.encdec_attention_layer_norm(
463
          self_attention_output + attention_output)
xinliupitt's avatar
xinliupitt committed
464
465
466
    if self._norm_first:
      source_attention_output = attention_output
      attention_output = self.output_layer_norm(attention_output)
467
468
469
470

    intermediate_output = self.intermediate_dense(attention_output)
    intermediate_output = self.intermediate_activation_layer(
        intermediate_output)
471
    intermediate_output = self._intermediate_dropout_layer(intermediate_output)
472
473
    layer_output = self.output_dense(intermediate_output)
    layer_output = self.output_dropout(layer_output)
xinliupitt's avatar
xinliupitt committed
474
475
476
477
    if self._norm_first:
      layer_output = source_attention_output + layer_output
    else:
      layer_output = self.output_layer_norm(layer_output + attention_output)
478
    return layer_output, cache