google / trax / 59

pending completion

Build # 59

Build Type

Pull #528

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/0f10bbfb57cfe5ebdcf71ca24a089a1ab688d46e">0f10bbfb5">&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/0f10bbfb57cfe5ebdcf71ca24a089a1ab688d46e">0f10bbfb5</a><a href="https://github.com/google/trax/commit/0f10bbfb57cfe5ebdcf71ca24a089a1ab688d46e">&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/0f10bbfb57cfe5ebdcf71ca24a089a1ab688d46e&quot;&gt;0f10bbfb5&lt;/a&gt;&lt;a href=&quot;https://github.com/google/trax/commit/0f10bbfb57cfe5ebdcf71ca24a089a1ab688d46e&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 #528: Refactor PureAttention and CausalAttention to highlight similarities/diffs.

Run Details

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

2681 of 10747 relevant lines covered (24.95%)

0.25 hits per line

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

32.8

/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 math
from trax.layers import base
from trax.layers import combinators as cb
from trax.layers import core
from trax.layers.base import Fn
from trax.math import numpy as jnp


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


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):
  def f(x):
    return jnp.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))
  return Fn(f'PaddingMask({pad})', f)


def EncoderDecoderMask():
  """Makes encoder-decoder mask from decoder input and a padding mask."""
  def f(decoder_input, padding_mask):
    padding_mask = jnp.reshape(
        padding_mask, (padding_mask.shape[0], 1, 1, padding_mask.shape[-1]))
    # Final mask shape is [batch, 1 for heads, decoder-len, encoder-len].
    return padding_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_with_state(self, inputs, weights=base.EMPTY_WEIGHTS,
                         state=base.EMPTY_STATE, rng=None):
    if self._mode != 'predict':
      x = inputs
      symbol_size = jnp.shape(x)[1]
      px = weights[:, :symbol_size, :]
      if self._dropout == 0:
        return (x + px, state)
      else:
        noise_shape = list(px.shape)
        for dim in self._dropout_broadcast_dims:
          noise_shape[dim] = 1
        keep_prob = 1.0 - self._dropout
        if math.backend_name() == 'jax':
          keep_prob = jax.lax.tie_in(x, jnp.full((), keep_prob, dtype=x.dtype))
        keep = math.random.bernoulli(rng, keep_prob, tuple(noise_shape))
        multiplier = keep.astype(x.dtype) / keep_prob
        return (x + px * multiplier, state)
    else:
      assert self._dropout == 0
      # 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.
      if inputs.shape[1] == 1:
        return (inputs + jnp.expand_dims(weights[0, state, :], 1), state + 1)
      else:
        emb = []
        for i in range(inputs.shape[0]):
          emb.append(jax.lax.dynamic_slice_in_dim(
              weights[0], state[i], inputs.shape[1], axis=0))
        return inputs + jnp.stack(emb, 0), state + inputs.shape[1]

  def new_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]
    weights = jnp.array(pe)  # Trainable parameters, initialized above.
    if self._mode == 'predict':
      batch_size = input_signature.shape[0]
      state = jnp.zeros((batch_size,), dtype=jnp.int32)
    else:
      state = base.EMPTY_STATE
    return weights, state


def DotProductAttention(query, key, value, mask, dropout, mode, rng):
  """Core dot product self-attention.

  Args:
    query: array of representations
    key: array of representations
    value: array of representations
    mask: attention-mask, gates attention
    dropout: float: dropout rate
    mode: 'eval' or 'train': whether to use dropout
    rng: Single-use random number generator (JAX PRNG key).

  Returns:
    Self attention for q, k, v arrays.
  """
  depth = jnp.shape(query)[-1]
  dots = jnp.matmul(query, jnp.swapaxes(key, -1, -2)) / jnp.sqrt(depth)
  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 math.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 - math.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 = math.random.bernoulli(rng, 1.0 - dropout, dots.shape)
    dots = jnp.where(keep, dots / (1.0 - dropout), jnp.zeros_like(dots))
  out = jnp.matmul(dots, value)
  return out


class PureAttention(base.Layer):
  """Layer constructor function for a pure attention layer."""

  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_with_state(self, x, weights, state, rng):
    """Pure transformer-style multi-headed attention.

    Args:
      x: inputs (q, k, v, mask)
      weights: parameters (none)
      state: parameters (none)
      rng: Single-use random number generator (JAX PRNG key).

    Returns:
      Pure Multi-headed attention result, and the mask.
    """
    del weights
    n_heads, dropout, mode = self._n_heads, self._dropout, self._mode
    q, k, v, mask = x
    batch_size = q.shape[0]
    d_feature = q.shape[-1]

    assert d_feature % n_heads == 0
    d_head = d_feature // n_heads

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

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

    res = _merge_heads(
        DotProductAttention(_split_into_heads(q),
                            _split_into_heads(k),
                            _split_into_heads(v),
                            mask,
                            dropout=dropout,
                            mode=mode,
                            rng=rng))
    return (res, mask), state  # Keep the mask.


def AttentionQKV(d_feature, n_heads=1, dropout=0.0, mode='train'):
  """Transformer-style multi-headed attention.

  Accepts inputs of the form q, k, v, mask.

  Args:
    d_feature: int:  dimensionality of feature embedding
    n_heads: int: number of attention heads
    dropout: float: dropout rate
    mode: str: 'train' or 'eval'

  Returns:
    Multi-headed self-attention result and the mask.
  """
  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),
  )


def Attention(d_feature, n_heads=1, dropout=0.0, mode='train'):
  """Transformer-style multi-headed attention.

  Accepts inputs of the form (x, mask) and constructs (q, k, v) from x.

  Args:
    d_feature: int:  dimensionality of feature embedding
    n_heads: int: number of attention heads
    dropout: float: dropout rate
    mode: str: 'train' or 'eval'

  Returns:
    Multi-headed self-attention result and the mask.
  """
  return cb.Serial(
      cb.Dup(), cb.Dup(),
      AttentionQKV(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
  )


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()
    depth = shape[-1]
    return jnp.zeros((batch_size, buffer_length, depth), 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."""
  assert math.backend_name() == 'jax', (
      'JAX backend is required to use the predict mode.')
  for x in inputs:
    assert x.shape[1] == 1, (
        'In predict mode the input sequence must be of length 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)


class DotProductCausalAttention(base.Layer):
  """A standard (non-memory-efficient) dot product attention implementation."""

  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_with_state(self, inputs, weights=base.EMPTY_WEIGHTS,
                         state=base.EMPTY_STATE, rng=None):
    del weights
    q, k, v = inputs
    if self._mode != 'predict':
      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 math.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)
    else:
      assert self._mode == 'predict'
      state = _fast_inference_update_state(inputs, state)
      (k, v, mask, _) = state

    res = DotProductAttention(
        q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
    return res, state

  def new_weights_and_state(self, input_signature):
    if self._mode != 'predict':
      return base.EMPTY_WEIGHTS, base.EMPTY_STATE

    assert self._mode == 'predict'
    weights = base.EMPTY_WEIGHTS
    # Buffer length is hardcoded for now. TODO(pkozakowski): Pass it from the
    # model.
    max_len = 2048
    state = _fast_inference_init_state(input_signature, max_len)
    return weights, state


def CausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):
  """Transformer-style multi-headed causal attention.

  Args:
    d_feature: int:  dimensionality of feature embedding
    n_heads: int: number of attention heads
    dropout: float: attention dropout
    mode: str: 'train' or 'eval'

  Returns:
    Multi-headed self-attention result.
  """
  assert d_feature % n_heads == 0
  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 = jnp.reshape(x, (batch_size, seq_len, n_heads, d_head))
      x = jnp.transpose(x, (0, 2, 1, 3))
      x = jnp.reshape(x, (-1, seq_len, d_head))
      return x
    return Fn('SplitIntoHeads', f)

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

      # (b_size*n_heads, seq_len, d_head) --> (b_size, seq_len, d_feature)
      x = jnp.reshape(x, (-1, n_heads, seq_len, d_head))
      x = jnp.transpose(x, (0, 2, 1, 3))
      x = jnp.reshape(x, (-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),
  )

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 math	1✔
23	from trax.layers import base	1✔
24	from trax.layers import combinators as cb	1✔
25	from trax.layers import core	1✔
26	from trax.layers.base import Fn	1✔
27	from trax.math import numpy as jnp	1✔
28
29
30	# Layers are always CamelCase, but functions in general are snake_case
31	# pylint: disable=invalid-name
32
33
34	def zero_pad(x, pad, axis):	1✔
35	"""Helper for jnp.pad with 0s for single-axis case."""
36	pad_widths = [(0, 0)] * len(x.shape)	1✔
37	pad_widths[axis] = pad # Padding on axis.	1✔
38	return jnp.pad(x, pad_widths, mode='constant',	1✔
39	constant_values=x.dtype.type(0))
40
41
42	def ShiftRight(n_shifts=1, mode='train'):	1✔
43	"""Layer to shift the tensor to the right by padding on axis 1."""
44	def f(x):	1✔
45	if mode == 'predict':	1✔
46	# Do nothing in predict mode, as then the sequence length is 1.
47	return x	×
48	padded = zero_pad(x, (n_shifts, 0), 1)	1✔
49	return padded[:, :-n_shifts]	1✔
50	return Fn(f'ShiftRight({n_shifts})', f)	1✔
51
52
53	def PaddingMask(pad=0):	1✔
54	def f(x):	×
55	return jnp.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))	×
56	return Fn(f'PaddingMask({pad})', f)	×
57
58
59	def EncoderDecoderMask():	1✔
60	"""Makes encoder-decoder mask from decoder input and a padding mask."""
61	def f(decoder_input, padding_mask):	×
62	padding_mask = jnp.reshape(	×
63	padding_mask, (padding_mask.shape[0], 1, 1, padding_mask.shape[-1]))
64	# Final mask shape is [batch, 1 for heads, decoder-len, encoder-len].
65	return padding_mask + jnp.zeros((1, 1, decoder_input.shape[1], 1))	×
66	return Fn('EncoderDecoderMask', f)	×
67
68
69	class PositionalEncoding(base.Layer):	1✔
70	"""Implements bare positional encoding."""
71
72	def __init__(self, max_len=2048, dropout=0.0, dropout_broadcast_dims=(-2,),	1✔
73	mode='train'):
74	super(PositionalEncoding, self).__init__()	1✔
75	self._max_len = max_len	1✔
76	if dropout >= 1.0:	1✔
77	raise ValueError('Dropout rates must be lower than 1.')	×
78	if mode == 'train':	1✔
79	self._dropout = dropout	1✔
80	else:
81	self._dropout = 0.0	×
82	self._dropout_broadcast_dims = dropout_broadcast_dims	1✔
83	self._mode = mode	1✔
84
85	def forward_with_state(self, inputs, weights=base.EMPTY_WEIGHTS,	1✔
86	state=base.EMPTY_STATE, rng=None):
87	if self._mode != 'predict':	1✔
88	x = inputs	1✔
89	symbol_size = jnp.shape(x)[1]	1✔
90	px = weights[:, :symbol_size, :]	1✔
91	if self._dropout == 0:	1✔
92	return (x + px, state)	1✔
93	else:
94	noise_shape = list(px.shape)	×
95	for dim in self._dropout_broadcast_dims:	×
96	noise_shape[dim] = 1	×
97	keep_prob = 1.0 - self._dropout	×
98	if math.backend_name() == 'jax':	×
99	keep_prob = jax.lax.tie_in(x, jnp.full((), keep_prob, dtype=x.dtype))	×
100	keep = math.random.bernoulli(rng, keep_prob, tuple(noise_shape))	×
101	multiplier = keep.astype(x.dtype) / keep_prob	×
102	return (x + px * multiplier, state)	×
103	else:
104	assert self._dropout == 0	×
105	# State in this class is only used for fast inference. In that case,
106	# the model is called with consecutive elements position-by-position.
107	# This positional encoding layer needs to store the index of the current
108	# position then and increment it on each call -- that's how state is used
109	# and updated below.
110	if inputs.shape[1] == 1:	×
111	return (inputs + jnp.expand_dims(weights[0, state, :], 1), state + 1)	×
112	else:
113	emb = []	×
114	for i in range(inputs.shape[0]):	×
115	emb.append(jax.lax.dynamic_slice_in_dim(	×
116	weights[0], state[i], inputs.shape[1], axis=0))
117	return inputs + jnp.stack(emb, 0), state + inputs.shape[1]	×
118
119	def new_weights_and_state(self, input_signature):	1✔
120	d_feature = input_signature.shape[-1]	1✔
121	pe = np.zeros((self._max_len, d_feature), dtype=np.float32)	1✔
122	position = np.arange(0, self._max_len)[:, np.newaxis]	1✔
123	div_term = np.exp(	1✔
124	np.arange(0, d_feature, 2) * -(np.log(10000.0) / d_feature))
125	pe[:, 0::2] = np.sin(position * div_term)	1✔
126	pe[:, 1::2] = np.cos(position * div_term)	1✔
127	pe = pe[np.newaxis, :, :] # [1, self._max_len, d_feature]	1✔
128	weights = jnp.array(pe) # Trainable parameters, initialized above.	1✔
129	if self._mode == 'predict':	1✔
130	batch_size = input_signature.shape[0]	×
131	state = jnp.zeros((batch_size,), dtype=jnp.int32)	×
132	else:
133	state = base.EMPTY_STATE	1✔
134	return weights, state	1✔
135
136
137	def DotProductAttention(query, key, value, mask, dropout, mode, rng):	1✔
138	"""Core dot product self-attention.
139
140	Args:
141	query: array of representations
142	key: array of representations
143	value: array of representations
144	mask: attention-mask, gates attention
145	dropout: float: dropout rate
146	mode: 'eval' or 'train': whether to use dropout
147	rng: Single-use random number generator (JAX PRNG key).
148
149	Returns:
150	Self attention for q, k, v arrays.
151	"""
152	depth = jnp.shape(query)[-1]	×
153	dots = jnp.matmul(query, jnp.swapaxes(key, -1, -2)) / jnp.sqrt(depth)	×
154	if mask is not None:	×
155	# TODO(kitaev): workaround for https://github.com/google/jax/issues/850
156	# We must ensure that both mask and the -1e9 constant have a data dependency
157	# on the input. Broadcasted copies of these use a lot of memory, so they
158	# should be computed at runtime (rather than being global constants).
159	if math.backend_name() == 'jax':	×
160	mask = jax.lax.tie_in(dots, mask)	×
161	# JAX's `full_like` already ties in -1e9 to dots.
162	dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))	×
163	# Softmax.
164	dots = jnp.exp(dots - math.logsumexp(dots, axis=-1, keepdims=True))	×
165	if dropout >= 1.0:	×
166	raise ValueError('Dropout rates must be lower than 1.')	×
167	if dropout is not None and dropout > 0.0 and mode == 'train':	×
168	keep = math.random.bernoulli(rng, 1.0 - dropout, dots.shape)	×
169	dots = jnp.where(keep, dots / (1.0 - dropout), jnp.zeros_like(dots))	×
170	out = jnp.matmul(dots, value)	×
171	return out	×
172
173
174	class PureAttention(base.Layer):	1✔
175	"""Layer constructor function for a pure attention layer."""
176
177	def __init__(self, n_heads=1, dropout=0.0, mode='train'):	1✔
178	super(PureAttention, self).__init__(n_in=4, n_out=2)	×
179	self._n_heads = n_heads	×
180	self._dropout = dropout	×
181	self._mode = mode	×
182
183	def forward_with_state(self, x, weights, state, rng):	1✔
184	"""Pure transformer-style multi-headed attention.
185
186	Args:
187	x: inputs (q, k, v, mask)
188	weights: parameters (none)
189	state: parameters (none)
190	rng: Single-use random number generator (JAX PRNG key).
191
192	Returns:
193	Pure Multi-headed attention result, and the mask.
194	"""
195	del weights	×
196	n_heads, dropout, mode = self._n_heads, self._dropout, self._mode	×
197	q, k, v, mask = x	×
198	batch_size = q.shape[0]	×
199	d_feature = q.shape[-1]	×
200
201	assert d_feature % n_heads == 0	×
202	d_head = d_feature // n_heads	×
203
204	def _split_into_heads(x):	×
205	"""Reshapes tensors to prepare for multi-headed computation."""
206	# (b_size, seqlen, d_feature) --> (b_size, n_heads, seqlen, d_head)
207	x = jnp.reshape(x, (batch_size, -1, n_heads, d_head))	×
208	x = jnp.transpose(x, (0, 2, 1, 3))	×
209	return x	×
210
211	def _merge_heads(x):	×
212	"""Undoes splitting, post multi-headed computation."""
213	# (b_size, n_heads, seqlen, d_head) --> (b_size, seqlen, d_feature)
214	x = jnp.transpose(x, (0, 2, 1, 3))	×
215	x = jnp.reshape(x, (batch_size, -1, n_heads * d_head))	×
216	return x	×
217
218	res = _merge_heads(	×
219	DotProductAttention(_split_into_heads(q),
220	_split_into_heads(k),
221	_split_into_heads(v),
222	mask,
223	dropout=dropout,
224	mode=mode,
225	rng=rng))
226	return (res, mask), state # Keep the mask.	×
227
228
229	def AttentionQKV(d_feature, n_heads=1, dropout=0.0, mode='train'):	1✔
230	"""Transformer-style multi-headed attention.
231
232	Accepts inputs of the form q, k, v, mask.
233
234	Args:
235	d_feature: int: dimensionality of feature embedding
236	n_heads: int: number of attention heads
237	dropout: float: dropout rate
238	mode: str: 'train' or 'eval'
239
240	Returns:
241	Multi-headed self-attention result and the mask.
242	"""
243	return cb.Serial(	×
244	cb.Parallel(
245	core.Dense(d_feature),
246	core.Dense(d_feature),
247	core.Dense(d_feature),
248	),
249	PureAttention( # pylint: disable=no-value-for-parameter
250	n_heads=n_heads, dropout=dropout, mode=mode),
251	core.Dense(d_feature),
252	)
253
254
255	def Attention(d_feature, n_heads=1, dropout=0.0, mode='train'):	1✔
256	"""Transformer-style multi-headed attention.
257
258	Accepts inputs of the form (x, mask) and constructs (q, k, v) from x.
259
260	Args:
261	d_feature: int: dimensionality of feature embedding
262	n_heads: int: number of attention heads
263	dropout: float: dropout rate
264	mode: str: 'train' or 'eval'
265
266	Returns:
267	Multi-headed self-attention result and the mask.
268	"""
269	return cb.Serial(	×
270	cb.Dup(), cb.Dup(),
271	AttentionQKV(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
272	)
273
274
275	def _fast_inference_init_state(input_signature, buffer_length):	1✔
276	"""Returns an initial state for causal attention layer fast inference."""
277	def zeros_for(batch_size, shape_dtype):	×
278	shape, dtype = shape_dtype.as_tuple()	×
279	depth = shape[-1]	×
280	return jnp.zeros((batch_size, buffer_length, depth), dtype=dtype)	×
281
282	batch_size = input_signature[0].shape[0]	×
283	k = zeros_for(batch_size, input_signature[1])	×
284	v = zeros_for(batch_size, input_signature[2])	×
285	mask = jnp.zeros((batch_size, 1, buffer_length))	×
286	seq_indices = jnp.zeros((batch_size,), dtype=jnp.int32)	×
287	return (k, v, mask, seq_indices)	×
288
289
290	def _fast_inference_update_state(inputs, state):	1✔
291	"""Updates state of a causal attention layer for fast inference."""
292	assert math.backend_name() == 'jax', (	×
293	'JAX backend is required to use the predict mode.')
294	for x in inputs:	×
295	assert x.shape[1] == 1, (	×
296	'In predict mode the input sequence must be of length 1.')
297	# Fast inference: run with only 1 query in each step, storing the sequence
298	# of keys and values calculated so far in state.
299	(_, new_k, new_v) = inputs	×
300	(ks, vs, mask, seq_indices) = state	×
301	batch_indices = jnp.arange(ks.shape[0])	×
302	ks = jax.ops.index_update(	×
303	ks, jax.ops.index[batch_indices, seq_indices, :], new_k[:, 0, :]
304	)
305	vs = jax.ops.index_update(	×
306	vs, jax.ops.index[batch_indices, seq_indices, :], new_v[:, 0, :]
307	)
308	mask = jax.ops.index_update(	×
309	mask, jax.ops.index[batch_indices, :, seq_indices], 1
310	)
311	return (ks, vs, mask, seq_indices + 1)	×
312
313
314	class DotProductCausalAttention(base.Layer):	1✔
315	"""A standard (non-memory-efficient) dot product attention implementation."""
316
317	def __init__(self, dropout=0.0, mode='train'):	1✔
318	super(DotProductCausalAttention, self).__init__(n_in=3, n_out=1)	×
319	self._dropout = dropout	×
320	self._mode = mode	×
321
322	def forward_with_state(self, inputs, weights=base.EMPTY_WEIGHTS,	1✔
323	state=base.EMPTY_STATE, rng=None):
324	del weights	×
325	q, k, v = inputs	×
326	if self._mode != 'predict':	×
327	mask_size = q.shape[-2]	×
328	# Not all backends define jnp.tril. However, using np.tril is inefficient
329	# in that it creates a large global constant. TODO(kitaev): try to find an
330	# alternative that works across all backends.
331	if math.backend_name() == 'jax':	×
332	mask = jnp.tril(	×
333	jnp.ones((1, mask_size, mask_size), dtype=np.bool_), k=0)
334	else:
335	mask = np.tril(	×
336	np.ones((1, mask_size, mask_size), dtype=np.bool_), k=0)
337	else:
338	assert self._mode == 'predict'	×
339	state = _fast_inference_update_state(inputs, state)	×
340	(k, v, mask, _) = state	×
341
342	res = DotProductAttention(	×
343	q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=rng)
344	return res, state	×
345
346	def new_weights_and_state(self, input_signature):	1✔
347	if self._mode != 'predict':	×
348	return base.EMPTY_WEIGHTS, base.EMPTY_STATE	×
349
350	assert self._mode == 'predict'	×
351	weights = base.EMPTY_WEIGHTS	×
352	# Buffer length is hardcoded for now. TODO(pkozakowski): Pass it from the
353	# model.
354	max_len = 2048	×
355	state = _fast_inference_init_state(input_signature, max_len)	×
356	return weights, state	×
357
358
359	def CausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):	1✔
360	"""Transformer-style multi-headed causal attention.
361
362	Args:
363	d_feature: int: dimensionality of feature embedding
364	n_heads: int: number of attention heads
365	dropout: float: attention dropout
366	mode: str: 'train' or 'eval'
367
368	Returns:
369	Multi-headed self-attention result.
370	"""
371	assert d_feature % n_heads == 0	×
372	d_head = d_feature // n_heads	×
373
374	def _split_into_heads():	×
375	"""Layer that reshapes tensors to prepare for multi-headed computation."""
376	def f(x):	×
377	batch_size = x.shape[0]	×
378	seq_len = x.shape[1]	×
379
380	# (b_size, seq_len, d_feature) --> (b_size*n_heads, seq_len, d_head)
381	x = jnp.reshape(x, (batch_size, seq_len, n_heads, d_head))	×
382	x = jnp.transpose(x, (0, 2, 1, 3))	×
383	x = jnp.reshape(x, (-1, seq_len, d_head))	×
384	return x	×
385	return Fn('SplitIntoHeads', f)	×
386
387	def _merge_heads():	×
388	"""Layer that undoes splitting, post multi-headed computation."""
389	def f(x):	×
390	seq_len = x.shape[1]	×
391
392	# (b_size*n_heads, seq_len, d_head) --> (b_size, seq_len, d_feature)
393	x = jnp.reshape(x, (-1, n_heads, seq_len, d_head))	×
394	x = jnp.transpose(x, (0, 2, 1, 3))	×
395	x = jnp.reshape(x, (-1, seq_len, n_heads * d_head))	×
396	return x	×
397	return Fn('MergeHeads', f)	×
398
399	return cb.Serial(	×
400	cb.Branch(
401	[core.Dense(d_feature), _split_into_heads()],
402	[core.Dense(d_feature), _split_into_heads()],
403	[core.Dense(d_feature), _split_into_heads()],
404	),
405	DotProductCausalAttention(dropout=dropout, mode=mode),
406	_merge_heads(),
407	core.Dense(d_feature),
408	)

google / trax / 59

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