2507

pending completion

Build # 2507

Build Type

Pull #827

travis-ci

Committed by

web-flow

Commit Message

<a href="https://github.com/google/trax/commit/<a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/17247f2693ef20bb8a626f19a47223202bc798b3">17247f269">&lt;a href=&quot;https://github.com/google/trax/commit/</a><a class="double-link" href="https://github.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/17247f2693ef20bb8a626f19a47223202bc798b3">17247f269</a><a href="https://github.com/google/trax/commit/17247f2693ef20bb8a626f19a47223202bc798b3">&quot;&gt;&amp;lt;a href=&amp;quot;https://github.com/google/trax/commit/&lt;/a&gt;&lt;a class=&quot;double-link&quot; href=&quot;https://github.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/17247f2693ef20bb8a626f19a47223202bc798b3&quot;&gt;17247f269&lt;/a&gt;&lt;a href=&quot;https://github.com/google/trax/commit/17247f2693ef20bb8a626f19a47223202bc798b3&quot;&gt;&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/google/trax/commit/&amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/google/trax/commit/&amp;amp;lt;a class=&amp;amp;quot;double-link&amp;amp;quot; href=&amp;amp;quot;https://git&amp;quot;&amp;gt;&amp;amp... (continued)

Pull Request Pull Request #827: Restore the attention visualizations from tensor2tensor.

Run Details

7 of 7 new or added lines in 1 file covered. (100.0%)

2623 of 9436 relevant lines covered (27.8%)

0.28 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

52.53

/trax/layers/attention.py

# coding=utf-8
# Copyright 2020 The Trax Authors.
#
# 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.

# Lint as: python3
"""Attention Layers."""

import jax
import numpy as np

from trax import fastmath
from trax.fastmath import numpy as jnp
from trax.layers import base
from trax.layers import combinators as cb
from trax.layers import core
from trax.layers.base import Fn


# Layers are always CamelCase, but functions in general are snake_case
# pylint: disable=invalid-name


def Attention(d_feature, n_heads=1, dropout=0.0, mode='train'):
  """Returns a layer that maps (activations, mask) to (new_activations, mask).

  This layer type represents one pass of multi-head self-attention, best
  known for its central role in Transformer models. Internally, it:

    - maps activations to `(queries, keys, values)` triples,
    - splits `queries`, `keys`, and `values` into multiple 'heads',
    - computes per-head attention weights from per-head `(queries, keys)`,
    - applies `mask` to screen out positions that come from padding tokens,
    - optionally applies dropout to attention weights,
    - uses attention weights to combine per-head `values` vectors, and
    - fuses per-head results into activations matching original input shapes.

  Args:
    d_feature: Depth/dimensionality of feature embedding.
    n_heads: Number of attention heads.
    dropout: Probababilistic rate for internal dropout applied to attention
        activations (based on query-key pairs) before dotting them with values.
    mode: Either 'train' or 'eval'.
  """
  return cb.Serial(
      cb.Dup(), cb.Dup(),  # TODO(jonni): replace with Select([0, 0, 0])
      AttentionQKV(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
  )


def AttentionQKV(d_feature, n_heads=1, dropout=0.0, mode='train'):
  """Returns a layer that maps (q, k, v, mask) to (activations, mask).

  See `Attention` above for further context/details.

  Args:
    d_feature: Depth/dimensionality of feature embedding.
    n_heads: Number of attention heads.
    dropout: Probababilistic rate for internal dropout applied to attention
        activations (based on query-key pairs) before dotting them with values.
    mode: Either 'train' or 'eval'.
  """
  return cb.Serial(
      cb.Parallel(
          core.Dense(d_feature),
          core.Dense(d_feature),
          core.Dense(d_feature),
      ),
      PureAttention(  # pylint: disable=no-value-for-parameter
          n_heads=n_heads, dropout=dropout, mode=mode),
      core.Dense(d_feature),
  )


class PureAttention(base.Layer):
  """Layer that maps from (queries, keys, values, mask) to (activations, mask).

  This layer type performs the inner workings of one pass of multi-head
  self-attention. It:

    - splits `queries`, `keys`, and `values` into multiple 'heads',
    - computes per-head attention weights from per-head `(queries, keys)`,
    - applies `mask` to screen out positions that come from padding tokens,
    - optionally applies dropout to attention weights,
    - uses attention weights to combine per-head `values` vectors, and
    - merges per-head results into activations matching original input shapes.
  """

  def __init__(self, n_heads=1, dropout=0.0, mode='train'):
    super(PureAttention, self).__init__(n_in=4, n_out=2)
    self._n_heads = n_heads
    self._dropout = dropout
    self._mode = mode

  def forward(self, inputs):
    """Returns attention-computed activations, unmodified mask, and state.

    Args:
      inputs: Tuple of (queries, keys, values, mask).

    Returns:
      The pair (activations, mask).
    """
    q, k, v, mask = inputs

    batch_size = q.shape[0]
    d_feature = q.shape[-1]
    n_heads = self._n_heads
    if d_feature % n_heads != 0:
      raise ValueError(
          f'Dimensionality of feature embedding ({d_feature}) is not a '
          f'multiple of the requested number of attention heads ({n_heads}).')

    d_head = d_feature // n_heads

    def _split_into_heads(x):
      """Reshapes tensors to prepare for multi-head computation."""
      # (b_size, seq_len, d_feature) --> (b_size, n_heads, seq_len, d_head)
      x = x.reshape((batch_size, -1, n_heads, d_head))
      x = x.transpose((0, 2, 1, 3))
      return x

    def _merge_heads(x):
      """Undoes splitting, after multi-head computation."""
      # (b_size, n_heads, seq_len, d_head) --> (b_size, seq_len, d_feature)
      x = x.transpose((0, 2, 1, 3))
      x = x.reshape((batch_size, -1, n_heads * d_head))
      return x

    per_head_results, dots = DotProductAttention(
        _split_into_heads(q),
        _split_into_heads(k),
        _split_into_heads(v),
        mask,
        dropout=self._dropout,
        mode=self._mode,
        rng=self.rng)
    if self._mode == 'viz':
      self.state = dots
    merged_results = _merge_heads(per_head_results)
    return (merged_results, mask)


def DotProductAttention(queries, keys, values, mask, dropout, mode, rng):
  """Computes new activations via masked attention-weighted sum of values.

  This function is the core of the attention mechanism. It:
    - computes per-head attention weights from per-head `(queries, keys)`,
    - applies `mask` to screen out positions that come from padding tokens,
    - optionally applies dropout to attention weights, and
    - uses attention weights to combine per-head `values` vectors.

  Args:
    queries: Per-head activations representing attention queries.
    keys: Per-head activations representing attention keys.
    values: Per-head activations to be combined by computed attention weights.
    mask: Mask that distinguishes positions with real content vs. padding.
    dropout: Probababilistic rate for dropout applied to attention activations
        (based on query-key pairs) before dotting them with values.
    mode: Either 'train' or eval'. Dropout applies only in 'train' mode.
    rng: Single-use random number generator (JAX PRNG key).

  Returns:
    Per-head activations resulting from masked per-head attention-weighted
    sum of per-head values.
  """
  d_feature = queries.shape[-1]
  dots = jnp.matmul(queries, jnp.swapaxes(keys, -1, -2)) / jnp.sqrt(d_feature)
  if mask is not None:
    # TODO(kitaev): workaround for https://github.com/google/jax/issues/850
    # We must ensure that both mask and the -1e9 constant have a data dependency
    # on the input. Broadcasted copies of these use a lot of memory, so they
    # should be computed at runtime (rather than being global constants).
    if fastmath.backend_name() == 'jax':
      mask = jax.lax.tie_in(dots, mask)
    # JAX's `full_like` already ties in -1e9 to dots.
    dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))
  # Softmax.
  dots = jnp.exp(dots - fastmath.logsumexp(dots, axis=-1, keepdims=True))
  if dropout >= 1.0:
    raise ValueError('Dropout rates must be lower than 1.')
  if dropout is not None and dropout > 0.0 and mode == 'train':
    keep = fastmath.random.bernoulli(rng, 1.0 - dropout, dots.shape)
    dots = jnp.where(keep, dots / (1.0 - dropout), jnp.zeros_like(dots))
  out = jnp.matmul(dots, values)
  return out, dots


def CausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):
  """Returns a layer that maps activations to activations, with causal masking.

  Like `Attention`, this layer type represents one pass of multi-head
  self-attention, but with causal masking in place of padding-based masking.

  Args:
    d_feature: Depth/dimensionality of feature embedding.
    n_heads: Number of attention heads.
    dropout: Probababilistic rate for internal dropout applied to attention
        activations (based on query-key pairs) before dotting them with values.
    mode: Either 'train' or 'eval'.
  """
  if d_feature % n_heads != 0:
    raise ValueError(
        f'Dimensionality of feature embedding ({d_feature}) is not a multiple '
        f'of the requested number of attention heads ({n_heads}).')

  d_head = d_feature // n_heads

  def _split_into_heads():
    """Layer that reshapes tensors to prepare for multi-headed computation."""
    def f(x):
      batch_size = x.shape[0]
      seq_len = x.shape[1]

      # (b_size, seq_len, d_feature) --> (b_size*n_heads, seq_len, d_head)
      x = x.reshape((batch_size, seq_len, n_heads, d_head))
      x = x.transpose((0, 2, 1, 3))
      x = x.reshape((-1, seq_len, d_head))
      return x
    return Fn('SplitIntoHeads', f)

  def _merge_heads():
    """Layer that undoes splitting, after multi-head computation."""
    def f(x):
      seq_len = x.shape[1]

      # (b_size*n_heads, seq_len, d_head) --> (b_size, seq_len, d_feature)
      x = x.reshape((-1, n_heads, seq_len, d_head))
      x = x.transpose((0, 2, 1, 3))
      x = x.reshape((-1, seq_len, n_heads * d_head))
      return x
    return Fn('MergeHeads', f)

  return cb.Serial(
      cb.Branch(
          [core.Dense(d_feature), _split_into_heads()],
          [core.Dense(d_feature), _split_into_heads()],
          [core.Dense(d_feature), _split_into_heads()],
      ),
      DotProductCausalAttention(dropout=dropout, mode=mode),
      _merge_heads(),
      core.Dense(d_feature),
  )


class DotProductCausalAttention(base.Layer):
  """Computes new activations via causally masked attention-weighted values."""

  def __init__(self, dropout=0.0, mode='train'):
    super(DotProductCausalAttention, self).__init__(n_in=3, n_out=1)
    self._dropout = dropout
    self._mode = mode

  def forward(self, inputs):
    q, k, v = inputs

    if self._mode == 'predict':
      self.state = _fast_inference_update_state(inputs, self.state)
      (k, v, mask, _) = self.state
    else:
      mask_size = q.shape[-2]
      # Not all backends define jnp.tril. However, using np.tril is inefficient
      # in that it creates a large global constant. TODO(kitaev): try to find an
      # alternative that works across all backends.
      if fastmath.backend_name() == 'jax':
        mask = jnp.tril(
            jnp.ones((1, mask_size, mask_size), dtype=np.bool_), k=0)
      else:
        mask = np.tril(
            np.ones((1, mask_size, mask_size), dtype=np.bool_), k=0)

    res, dots = DotProductAttention(
        q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=self.rng)
    if self._mode == 'viz':
      self.state = dots
    return res

  def init_weights_and_state(self, input_signature):
    if self._mode == 'predict':
      max_len = 2048  # Hardcoded.  TODO(pkozakowski): Pass it from the model.
      self.state = _fast_inference_init_state(input_signature, max_len)


def zero_pad(x, pad, axis):
  """Helper for jnp.pad with 0s for single-axis case."""
  pad_widths = [(0, 0)] * len(x.shape)
  pad_widths[axis] = pad  # Padding on axis.
  return jnp.pad(x, pad_widths, mode='constant',
                 constant_values=x.dtype.type(0))


def ShiftRight(n_shifts=1, mode='train'):
  """Layer to shift the tensor to the right by padding on axis 1."""
  def f(x):
    if mode == 'predict':
      # Do nothing in predict mode, as then the sequence length is 1.
      return x
    padded = zero_pad(x, (n_shifts, 0), 1)
    return padded[:, :-n_shifts]
  return Fn(f'ShiftRight({n_shifts})', f)


def PaddingMask(pad=0):
  """Layer to distinguish positions with real content/tokens vs. padding."""
  def f(x):
    # TODO(jonni): Check/require that len(x.shape) == 2?
    batch_size = x.shape[0]
    d_feature = x.shape[-1]
    content_positions = (x != pad)
    return content_positions.reshape((batch_size, 1, 1, d_feature))
  return Fn(f'PaddingMask({pad})', f)


def EncoderDecoderMask():
  """Makes encoder-decoder mask from decoder input and a padding mask."""
  def f(decoder_input, mask):
    batch_size = mask.shape[0]
    d_feature = mask.shape[-1]
    mask = mask.reshape((batch_size, 1, 1, d_feature))
    # Final mask shape is [batch, 1 for heads, decoder-len, encoder-len].
    return mask + jnp.zeros((1, 1, decoder_input.shape[1], 1))
  return Fn('EncoderDecoderMask', f)


class PositionalEncoding(base.Layer):
  """Implements bare positional encoding."""

  def __init__(self, max_len=2048, dropout=0.0, dropout_broadcast_dims=(-2,),
               mode='train'):
    super(PositionalEncoding, self).__init__()
    self._max_len = max_len
    if dropout >= 1.0:
      raise ValueError('Dropout rates must be lower than 1.')
    if mode == 'train':
      self._dropout = dropout
    else:
      self._dropout = 0.0
    self._dropout_broadcast_dims = dropout_broadcast_dims
    self._mode = mode

  def forward(self, inputs):
    if self._mode != 'predict':
      x = inputs
      symbol_size = jnp.shape(x)[1]
      px = self.weights[:, :symbol_size, :]
      if self._dropout == 0:
        return x + px
      else:
        noise_shape = list(px.shape)
        for dim in self._dropout_broadcast_dims:
          noise_shape[dim] = 1
        keep_prob = 1.0 - self._dropout
        if fastmath.backend_name() == 'jax':
          keep_prob = jax.lax.tie_in(x, jnp.full((), keep_prob, dtype=x.dtype))
        keep = fastmath.random.bernoulli(self.rng, keep_prob,
                                         tuple(noise_shape))
        multiplier = keep.astype(x.dtype) / keep_prob
        return x + px * multiplier
    else:
      if self._dropout != 0:
        raise ValueError(f'In predict mode, but dropout rate '
                         f'({self._dropout}) is not zero.')

      # State in this class is only used for fast inference. In that case,
      # the model is called with consecutive elements position-by-position.
      # This positional encoding layer needs to store the index of the current
      # position then and increment it on each call -- that's how state is used
      # and updated below.
      state = self.state
      if inputs.shape[1] == 1:
        self.state = state + 1
        return inputs + jnp.expand_dims(self.weights[0, state, :], 1)
      else:
        emb = []
        for i in range(inputs.shape[0]):
          emb.append(jax.lax.dynamic_slice_in_dim(
              self.weights[0], state[i], inputs.shape[1], axis=0))
        self.state = state + inputs.shape[1]
        return inputs + jnp.stack(emb, 0)

  def init_weights_and_state(self, input_signature):
    d_feature = input_signature.shape[-1]
    pe = np.zeros((self._max_len, d_feature), dtype=np.float32)
    position = np.arange(0, self._max_len)[:, np.newaxis]
    div_term = np.exp(
        np.arange(0, d_feature, 2) * -(np.log(10000.0) / d_feature))
    pe[:, 0::2] = np.sin(position * div_term)
    pe[:, 1::2] = np.cos(position * div_term)
    pe = pe[np.newaxis, :, :]  # [1, self._max_len, d_feature]
    self.weights = jnp.array(pe)  # Trainable parameters, initialized above.
    if self._mode == 'predict':
      batch_size = input_signature.shape[0]
      self.state = jnp.zeros((batch_size,), dtype=jnp.int32)


def _fast_inference_init_state(input_signature, buffer_length):
  """Returns an initial state for causal attention layer fast inference."""
  def zeros_for(batch_size, shape_dtype):
    shape, dtype = shape_dtype.as_tuple()
    d_feature = shape[-1]
    return jnp.zeros((batch_size, buffer_length, d_feature), dtype=dtype)

  batch_size = input_signature[0].shape[0]
  k = zeros_for(batch_size, input_signature[1])
  v = zeros_for(batch_size, input_signature[2])
  mask = jnp.zeros((batch_size, 1, buffer_length))
  seq_indices = jnp.zeros((batch_size,), dtype=jnp.int32)
  return (k, v, mask, seq_indices)


def _fast_inference_update_state(inputs, state):
  """Updates state of a causal attention layer for fast inference."""
  if fastmath.backend_name() != 'jax':
    raise ValueError(f'JAX backend is required in predict mode, but found '
                     f'backend ({fastmath.backend_nameO()}).')
  for x in inputs:
    if x.shape[1] != 1:
      raise ValueError(f'In predict mode, input sequence must have length 1, '
                       f'instead has length {x.shape[1]}.')
  # Fast inference: run with only 1 query in each step, storing the sequence
  # of keys and values calculated so far in state.
  (_, new_k, new_v) = inputs
  (ks, vs, mask, seq_indices) = state
  batch_indices = jnp.arange(ks.shape[0])
  ks = jax.ops.index_update(
      ks, jax.ops.index[batch_indices, seq_indices, :], new_k[:, 0, :])
  vs = jax.ops.index_update(
      vs, jax.ops.index[batch_indices, seq_indices, :], new_v[:, 0, :])
  mask = jax.ops.index_update(
      mask, jax.ops.index[batch_indices, :, seq_indices], 1)
  return (ks, vs, mask, seq_indices + 1)

1	# coding=utf-8
2	# Copyright 2020 The Trax Authors.
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	# http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15
16	# Lint as: python3
17	"""Attention Layers."""	1✔
18
19	import jax	1✔
20	import numpy as np	1✔
21
22	from trax import fastmath	1✔
23	from trax.fastmath import numpy as jnp	1✔
24	from trax.layers import base	1✔
25	from trax.layers import combinators as cb	1✔
26	from trax.layers import core	1✔
27	from trax.layers.base import Fn	1✔
28
29
30	# Layers are always CamelCase, but functions in general are snake_case
31	# pylint: disable=invalid-name
32
33
34	def Attention(d_feature, n_heads=1, dropout=0.0, mode='train'):	1✔
35	"""Returns a layer that maps (activations, mask) to (new_activations, mask).
36
37	This layer type represents one pass of multi-head self-attention, best
38	known for its central role in Transformer models. Internally, it:
39
40	- maps activations to `(queries, keys, values)` triples,
41	- splits `queries`, `keys`, and `values` into multiple 'heads',
42	- computes per-head attention weights from per-head `(queries, keys)`,
43	- applies `mask` to screen out positions that come from padding tokens,
44	- optionally applies dropout to attention weights,
45	- uses attention weights to combine per-head `values` vectors, and
46	- fuses per-head results into activations matching original input shapes.
47
48	Args:
49	d_feature: Depth/dimensionality of feature embedding.
50	n_heads: Number of attention heads.
51	dropout: Probababilistic rate for internal dropout applied to attention
52	activations (based on query-key pairs) before dotting them with values.
53	mode: Either 'train' or 'eval'.
54	"""
55	return cb.Serial(	×
56	cb.Dup(), cb.Dup(), # TODO(jonni): replace with Select([0, 0, 0])
57	AttentionQKV(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
58	)
59
60
61	def AttentionQKV(d_feature, n_heads=1, dropout=0.0, mode='train'):	1✔
62	"""Returns a layer that maps (q, k, v, mask) to (activations, mask).
63
64	See `Attention` above for further context/details.
65
66	Args:
67	d_feature: Depth/dimensionality of feature embedding.
68	n_heads: Number of attention heads.
69	dropout: Probababilistic rate for internal dropout applied to attention
70	activations (based on query-key pairs) before dotting them with values.
71	mode: Either 'train' or 'eval'.
72	"""
73	return cb.Serial(	×
74	cb.Parallel(
75	core.Dense(d_feature),
76	core.Dense(d_feature),
77	core.Dense(d_feature),
78	),
79	PureAttention( # pylint: disable=no-value-for-parameter
80	n_heads=n_heads, dropout=dropout, mode=mode),
81	core.Dense(d_feature),
82	)
83
84
85	class PureAttention(base.Layer):	1✔
86	"""Layer that maps from (queries, keys, values, mask) to (activations, mask).
87
88	This layer type performs the inner workings of one pass of multi-head
89	self-attention. It:
90
91	- splits `queries`, `keys`, and `values` into multiple 'heads',
92	- computes per-head attention weights from per-head `(queries, keys)`,
93	- applies `mask` to screen out positions that come from padding tokens,
94	- optionally applies dropout to attention weights,
95	- uses attention weights to combine per-head `values` vectors, and
96	- merges per-head results into activations matching original input shapes.
97	"""
98
99	def __init__(self, n_heads=1, dropout=0.0, mode='train'):	1✔
100	super(PureAttention, self).__init__(n_in=4, n_out=2)	×
101	self._n_heads = n_heads	×
102	self._dropout = dropout	×
103	self._mode = mode	×
104
105	def forward(self, inputs):	1✔
106	"""Returns attention-computed activations, unmodified mask, and state.
107
108	Args:
109	inputs: Tuple of (queries, keys, values, mask).
110
111	Returns:
112	The pair (activations, mask).
113	"""
114	q, k, v, mask = inputs	×
115
116	batch_size = q.shape[0]	×
117	d_feature = q.shape[-1]	×
118	n_heads = self._n_heads	×
119	if d_feature % n_heads != 0:	×
120	raise ValueError(	×
121	f'Dimensionality of feature embedding ({d_feature}) is not a '
122	f'multiple of the requested number of attention heads ({n_heads}).')
123
124	d_head = d_feature // n_heads	×
125
126	def _split_into_heads(x):	×
127	"""Reshapes tensors to prepare for multi-head computation."""
128	# (b_size, seq_len, d_feature) --> (b_size, n_heads, seq_len, d_head)
129	x = x.reshape((batch_size, -1, n_heads, d_head))	×
130	x = x.transpose((0, 2, 1, 3))	×
131	return x	×
132
133	def _merge_heads(x):	×
134	"""Undoes splitting, after multi-head computation."""
135	# (b_size, n_heads, seq_len, d_head) --> (b_size, seq_len, d_feature)
136	x = x.transpose((0, 2, 1, 3))	×
137	x = x.reshape((batch_size, -1, n_heads * d_head))	×
138	return x	×
139
140	per_head_results, dots = DotProductAttention(	×
141	_split_into_heads(q),
142	_split_into_heads(k),
143	_split_into_heads(v),
144	mask,
145	dropout=self._dropout,
146	mode=self._mode,
147	rng=self.rng)
148	if self._mode == 'viz':	×
149	self.state = dots	×
150	merged_results = _merge_heads(per_head_results)	×
151	return (merged_results, mask)	×
152
153
154	def DotProductAttention(queries, keys, values, mask, dropout, mode, rng):	1✔
155	"""Computes new activations via masked attention-weighted sum of values.
156
157	This function is the core of the attention mechanism. It:
158	- computes per-head attention weights from per-head `(queries, keys)`,
159	- applies `mask` to screen out positions that come from padding tokens,
160	- optionally applies dropout to attention weights, and
161	- uses attention weights to combine per-head `values` vectors.
162
163	Args:
164	queries: Per-head activations representing attention queries.
165	keys: Per-head activations representing attention keys.
166	values: Per-head activations to be combined by computed attention weights.
167	mask: Mask that distinguishes positions with real content vs. padding.
168	dropout: Probababilistic rate for dropout applied to attention activations
169	(based on query-key pairs) before dotting them with values.
170	mode: Either 'train' or eval'. Dropout applies only in 'train' mode.
171	rng: Single-use random number generator (JAX PRNG key).
172
173	Returns:
174	Per-head activations resulting from masked per-head attention-weighted
175	sum of per-head values.
176	"""
177	d_feature = queries.shape[-1]	1✔
178	dots = jnp.matmul(queries, jnp.swapaxes(keys, -1, -2)) / jnp.sqrt(d_feature)	1✔
179	if mask is not None:	1✔
180	# TODO(kitaev): workaround for https://github.com/google/jax/issues/850
181	# We must ensure that both mask and the -1e9 constant have a data dependency
182	# on the input. Broadcasted copies of these use a lot of memory, so they
183	# should be computed at runtime (rather than being global constants).
184	if fastmath.backend_name() == 'jax':	1✔
185	mask = jax.lax.tie_in(dots, mask)	×
186	# JAX's `full_like` already ties in -1e9 to dots.
187	dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))	1✔
188	# Softmax.
189	dots = jnp.exp(dots - fastmath.logsumexp(dots, axis=-1, keepdims=True))	1✔
190	if dropout >= 1.0:	1✔
191	raise ValueError('Dropout rates must be lower than 1.')	×
192	if dropout is not None and dropout > 0.0 and mode == 'train':	1✔
193	keep = fastmath.random.bernoulli(rng, 1.0 - dropout, dots.shape)	1✔
194	dots = jnp.where(keep, dots / (1.0 - dropout), jnp.zeros_like(dots))	1✔
195	out = jnp.matmul(dots, values)	1✔
196	return out, dots	1✔
197
198
199	def CausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):	1✔
200	"""Returns a layer that maps activations to activations, with causal masking.
201
202	Like `Attention`, this layer type represents one pass of multi-head
203	self-attention, but with causal masking in place of padding-based masking.
204
205	Args:
206	d_feature: Depth/dimensionality of feature embedding.
207	n_heads: Number of attention heads.
208	dropout: Probababilistic rate for internal dropout applied to attention
209	activations (based on query-key pairs) before dotting them with values.
210	mode: Either 'train' or 'eval'.
211	"""
212	if d_feature % n_heads != 0:	1✔
213	raise ValueError(	×
214	f'Dimensionality of feature embedding ({d_feature}) is not a multiple '
215	f'of the requested number of attention heads ({n_heads}).')
216
217	d_head = d_feature // n_heads	1✔
218
219	def _split_into_heads():	1✔
220	"""Layer that reshapes tensors to prepare for multi-headed computation."""
221	def f(x):	1✔
222	batch_size = x.shape[0]	1✔
223	seq_len = x.shape[1]	1✔
224
225	# (b_size, seq_len, d_feature) --> (b_size*n_heads, seq_len, d_head)
226	x = x.reshape((batch_size, seq_len, n_heads, d_head))	1✔
227	x = x.transpose((0, 2, 1, 3))	1✔
228	x = x.reshape((-1, seq_len, d_head))	1✔
229	return x	1✔
230	return Fn('SplitIntoHeads', f)	1✔
231
232	def _merge_heads():	1✔
233	"""Layer that undoes splitting, after multi-head computation."""
234	def f(x):	1✔
235	seq_len = x.shape[1]	1✔
236
237	# (b_size*n_heads, seq_len, d_head) --> (b_size, seq_len, d_feature)
238	x = x.reshape((-1, n_heads, seq_len, d_head))	1✔
239	x = x.transpose((0, 2, 1, 3))	1✔
240	x = x.reshape((-1, seq_len, n_heads * d_head))	1✔
241	return x	1✔
242	return Fn('MergeHeads', f)	1✔
243
244	return cb.Serial(	1✔
245	cb.Branch(
246	[core.Dense(d_feature), _split_into_heads()],
247	[core.Dense(d_feature), _split_into_heads()],
248	[core.Dense(d_feature), _split_into_heads()],
249	),
250	DotProductCausalAttention(dropout=dropout, mode=mode),
251	_merge_heads(),
252	core.Dense(d_feature),
253	)
254
255
256	class DotProductCausalAttention(base.Layer):	1✔
257	"""Computes new activations via causally masked attention-weighted values."""
258
259	def __init__(self, dropout=0.0, mode='train'):	1✔
260	super(DotProductCausalAttention, self).__init__(n_in=3, n_out=1)	1✔
261	self._dropout = dropout	1✔
262	self._mode = mode	1✔
263
264	def forward(self, inputs):	1✔
265	q, k, v = inputs	1✔
266
267	if self._mode == 'predict':	1✔
268	self.state = _fast_inference_update_state(inputs, self.state)	×
269	(k, v, mask, _) = self.state	×
270	else:
271	mask_size = q.shape[-2]	1✔
272	# Not all backends define jnp.tril. However, using np.tril is inefficient
273	# in that it creates a large global constant. TODO(kitaev): try to find an
274	# alternative that works across all backends.
275	if fastmath.backend_name() == 'jax':	1✔
276	mask = jnp.tril(	×
277	jnp.ones((1, mask_size, mask_size), dtype=np.bool_), k=0)
278	else:
279	mask = np.tril(	1✔
280	np.ones((1, mask_size, mask_size), dtype=np.bool_), k=0)
281
282	res, dots = DotProductAttention(	1✔
283	q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=self.rng)
284	if self._mode == 'viz':	1✔
285	self.state = dots	×
286	return res	1✔
287
288	def init_weights_and_state(self, input_signature):	1✔
289	if self._mode == 'predict':	1✔
290	max_len = 2048 # Hardcoded. TODO(pkozakowski): Pass it from the model.	×
291	self.state = _fast_inference_init_state(input_signature, max_len)	×
292
293
294	def zero_pad(x, pad, axis):	1✔
295	"""Helper for jnp.pad with 0s for single-axis case."""
296	pad_widths = [(0, 0)] * len(x.shape)	1✔
297	pad_widths[axis] = pad # Padding on axis.	1✔
298	return jnp.pad(x, pad_widths, mode='constant',	1✔
299	constant_values=x.dtype.type(0))
300
301
302	def ShiftRight(n_shifts=1, mode='train'):	1✔
303	"""Layer to shift the tensor to the right by padding on axis 1."""
304	def f(x):	1✔
305	if mode == 'predict':	1✔
306	# Do nothing in predict mode, as then the sequence length is 1.
307	return x	×
308	padded = zero_pad(x, (n_shifts, 0), 1)	1✔
309	return padded[:, :-n_shifts]	1✔
310	return Fn(f'ShiftRight({n_shifts})', f)	1✔
311
312
313	def PaddingMask(pad=0):	1✔
314	"""Layer to distinguish positions with real content/tokens vs. padding."""
315	def f(x):	×
316	# TODO(jonni): Check/require that len(x.shape) == 2?
317	batch_size = x.shape[0]	×
318	d_feature = x.shape[-1]	×
319	content_positions = (x != pad)	×
320	return content_positions.reshape((batch_size, 1, 1, d_feature))	×
321	return Fn(f'PaddingMask({pad})', f)	×
322
323
324	def EncoderDecoderMask():	1✔
325	"""Makes encoder-decoder mask from decoder input and a padding mask."""
326	def f(decoder_input, mask):	×
327	batch_size = mask.shape[0]	×
328	d_feature = mask.shape[-1]	×
329	mask = mask.reshape((batch_size, 1, 1, d_feature))	×
330	# Final mask shape is [batch, 1 for heads, decoder-len, encoder-len].
331	return mask + jnp.zeros((1, 1, decoder_input.shape[1], 1))	×
332	return Fn('EncoderDecoderMask', f)	×
333
334
335	class PositionalEncoding(base.Layer):	1✔
336	"""Implements bare positional encoding."""
337
338	def __init__(self, max_len=2048, dropout=0.0, dropout_broadcast_dims=(-2,),	1✔
339	mode='train'):
340	super(PositionalEncoding, self).__init__()	1✔
341	self._max_len = max_len	1✔
342	if dropout >= 1.0:	1✔
343	raise ValueError('Dropout rates must be lower than 1.')	×
344	if mode == 'train':	1✔
345	self._dropout = dropout	1✔
346	else:
347	self._dropout = 0.0	×
348	self._dropout_broadcast_dims = dropout_broadcast_dims	1✔
349	self._mode = mode	1✔
350
351	def forward(self, inputs):	1✔
352	if self._mode != 'predict':	1✔
353	x = inputs	1✔
354	symbol_size = jnp.shape(x)[1]	1✔
355	px = self.weights[:, :symbol_size, :]	1✔
356	if self._dropout == 0:	1✔
357	return x + px	1✔
358	else:
359	noise_shape = list(px.shape)	×
360	for dim in self._dropout_broadcast_dims:	×
361	noise_shape[dim] = 1	×
362	keep_prob = 1.0 - self._dropout	×
363	if fastmath.backend_name() == 'jax':	×
364	keep_prob = jax.lax.tie_in(x, jnp.full((), keep_prob, dtype=x.dtype))	×
365	keep = fastmath.random.bernoulli(self.rng, keep_prob,	×
366	tuple(noise_shape))
367	multiplier = keep.astype(x.dtype) / keep_prob	×
368	return x + px * multiplier	×
369	else:
370	if self._dropout != 0:	×
371	raise ValueError(f'In predict mode, but dropout rate '	×
372	f'({self._dropout}) is not zero.')
373
374	# State in this class is only used for fast inference. In that case,
375	# the model is called with consecutive elements position-by-position.
376	# This positional encoding layer needs to store the index of the current
377	# position then and increment it on each call -- that's how state is used
378	# and updated below.
379	state = self.state	×
380	if inputs.shape[1] == 1:	×
381	self.state = state + 1	×
382	return inputs + jnp.expand_dims(self.weights[0, state, :], 1)	×
383	else:
384	emb = []	×
385	for i in range(inputs.shape[0]):	×
386	emb.append(jax.lax.dynamic_slice_in_dim(	×
387	self.weights[0], state[i], inputs.shape[1], axis=0))
388	self.state = state + inputs.shape[1]	×
389	return inputs + jnp.stack(emb, 0)	×
390
391	def init_weights_and_state(self, input_signature):	1✔
392	d_feature = input_signature.shape[-1]	1✔
393	pe = np.zeros((self._max_len, d_feature), dtype=np.float32)	1✔
394	position = np.arange(0, self._max_len)[:, np.newaxis]	1✔
395	div_term = np.exp(	1✔
396	np.arange(0, d_feature, 2) * -(np.log(10000.0) / d_feature))
397	pe[:, 0::2] = np.sin(position * div_term)	1✔
398	pe[:, 1::2] = np.cos(position * div_term)	1✔
399	pe = pe[np.newaxis, :, :] # [1, self._max_len, d_feature]	1✔
400	self.weights = jnp.array(pe) # Trainable parameters, initialized above.	1✔
401	if self._mode == 'predict':	1✔
402	batch_size = input_signature.shape[0]	×
403	self.state = jnp.zeros((batch_size,), dtype=jnp.int32)	×
404
405
406	def _fast_inference_init_state(input_signature, buffer_length):	1✔
407	"""Returns an initial state for causal attention layer fast inference."""
408	def zeros_for(batch_size, shape_dtype):	×
409	shape, dtype = shape_dtype.as_tuple()	×
410	d_feature = shape[-1]	×
411	return jnp.zeros((batch_size, buffer_length, d_feature), dtype=dtype)	×
412
413	batch_size = input_signature[0].shape[0]	×
414	k = zeros_for(batch_size, input_signature[1])	×
415	v = zeros_for(batch_size, input_signature[2])	×
416	mask = jnp.zeros((batch_size, 1, buffer_length))	×
417	seq_indices = jnp.zeros((batch_size,), dtype=jnp.int32)	×
418	return (k, v, mask, seq_indices)	×
419
420
421	def _fast_inference_update_state(inputs, state):	1✔
422	"""Updates state of a causal attention layer for fast inference."""
423	if fastmath.backend_name() != 'jax':	×
424	raise ValueError(f'JAX backend is required in predict mode, but found '	×
425	f'backend ({fastmath.backend_nameO()}).')
426	for x in inputs:	×
427	if x.shape[1] != 1:	×
428	raise ValueError(f'In predict mode, input sequence must have length 1, '	×
429	f'instead has length {x.shape[1]}.')
430	# Fast inference: run with only 1 query in each step, storing the sequence
431	# of keys and values calculated so far in state.
432	(_, new_k, new_v) = inputs	×
433	(ks, vs, mask, seq_indices) = state	×
434	batch_indices = jnp.arange(ks.shape[0])	×
435	ks = jax.ops.index_update(	×
436	ks, jax.ops.index[batch_indices, seq_indices, :], new_k[:, 0, :])
437	vs = jax.ops.index_update(	×
438	vs, jax.ops.index[batch_indices, seq_indices, :], new_v[:, 0, :])
439	mask = jax.ops.index_update(	×
440	mask, jax.ops.index[batch_indices, :, seq_indices], 1)
441	return (ks, vs, mask, seq_indices + 1)	×

google / trax / 2507

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous