2733

pending completion

Build # 2733

Build Type

Pull #847

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/7ae800c7e65202335d3e74ff42d88345192fd02a">7ae800c7e">&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/7ae800c7e65202335d3e74ff42d88345192fd02a">7ae800c7e</a><a href="https://github.com/google/trax/commit/7ae800c7e65202335d3e74ff42d88345192fd02a">&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/7ae800c7e65202335d3e74ff42d88345192fd02a&quot;&gt;7ae800c7e&lt;/a&gt;&lt;a href=&quot;https://github.com/google/trax/commit/7ae800c7e65202335d3e74ff42d88345192fd02a&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 #847: Adds `scan` to tf_numpy/extensions.

Run Details

52 of 52 new or added lines in 3 files covered. (100.0%)

2618 of 9551 relevant lines covered (27.41%)

0.27 hits per line

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

61.47

/trax/fastmath/ops.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
"""Trax accelerated math operations for fast computing on GPUs and TPUs.

Import these operations directly from fastmath and import fastmath.numpy as np::

    from trax import fastmath
    from trax.fastmath import numpy as np

    x = np.array([1.0, 2.0])  # Use like numpy.
    y = np.exp(x)  # Common numpy ops are available and accelerated.
    z = fastmath.logsumexp(y)  # Special operations available from fastmath.

Trax uses either TensorFlow 2 or JAX as backend for accelerating operations.
You can select which one to use (e.g., for  debugging) with `use_backend`.
"""

import contextlib

import gin
from trax.fastmath.jax import JAX_BACKEND
from trax.fastmath.numpy import NUMPY_BACKEND
from trax.fastmath.tf import TF_BACKEND

# For numpy and random modules, we need to call "backend()" lazily, only when
# the function is called -- so that it can be set by gin configs.
# (Otherwise, backend() is called on import before gin-config is parsed.)
# To do that, we make objects to encapsulated these modules.


# A class that just forwards attribute accesses to backend's numpy object.
class NumpyBackend(object):
  """Numpy functions accelerated to run on GPUs and TPUs. Use like numpy."""

  def __getattr__(self, attr):
    return getattr(backend()['np'], attr)

numpy = NumpyBackend()


class RandomBackend(object):
  """Backend providing random functions."""

  def get_prng(self, seed):
    return backend()['random_get_prng'](seed)

  def split(self, prng, num=2):
    return backend()['random_split'](prng, num)

  def uniform(self, *args, **kwargs):
    return backend()['random_uniform'](*args, **kwargs)

  def randint(self, *args, **kwargs):
    return backend()['random_randint'](*args, **kwargs)

  def normal(self, *args, **kwargs):
    return backend()['random_normal'](*args, **kwargs)

  def bernoulli(self, *args, **kwargs):
    return backend()['random_bernoulli'](*args, **kwargs)


random = RandomBackend()


def logsumexp(*args, **kwargs):
  """Computes the log of the sum of exponentials of input elements."""
  return backend()['logsumexp'](*args, **kwargs)


def expit(*args, **kwargs):
  """Computes the expit (sigmoid) function."""
  return backend()['expit'](*args, **kwargs)


def sigmoid(*args, **kwargs):
  """Computes the sigmoid (expit) function."""
  return backend()['expit'](*args, **kwargs)


def erf(*args, **kwargs):
  """Computes the erf function."""
  return backend()['erf'](*args, **kwargs)


def conv(*args, **kwargs):
  """Computes a generalized convolution."""
  return backend()['conv'](*args, **kwargs)


def avg_pool(*args, **kwargs):
  """Average pooling."""
  return backend()['avg_pool'](*args, **kwargs)


def max_pool(*args, **kwargs):
  """Max pooling."""
  return backend()['max_pool'](*args, **kwargs)


def sum_pool(*args, **kwargs):
  """Sum pooling."""
  return backend()['sum_pool'](*args, **kwargs)


def sort_key_val(*args, **kwargs):
  """Sorts keys along dimension and applies same permutation to values."""
  return backend()['sort_key_val'](*args, **kwargs)


def scan(*args, **kwargs):
  """Scan to make recurrent functions run faster on accelerators."""
  return backend()['scan'](*args, **kwargs)


def remat(*args, **kwargs):
  """Recompute everything in the backward pass to same memory."""
  return backend()['remat'](*args, **kwargs)


def cond(*args, **kwargs):
  """Conditional computation to run on accelerators."""
  return backend()['cond'](*args, **kwargs)


def lt(*args, **kwargs):
  """Less-than function for backends that do not override <."""
  return backend()['lt'](*args, **kwargs)


def stop_gradient(*args, **kwargs):
  """Identity on the forward pass but 0 (no gradient) on the backward pass."""
  return backend()['stop_gradient'](*args, **kwargs)


_disable_jit = False


def jit(*args, **kwargs):
  """Just-In-Time compiles the given function for use on accelerators."""
  global _disable_jit
  if _disable_jit:
    return args[0]  # jit(f, **unused_now_jit_kwargs) = f
  return backend()['jit'](*args, **kwargs)


def disable_jit():
  """Disables JIT-compilation; helpful for debugging."""
  global _disable_jit
  _disable_jit = True


def grad(*args, **kwargs):
  """Computes the gradient of the specified function (returns a function)."""
  return backend()['grad'](*args, **kwargs)


def value_and_grad(*args, **kwargs):
  """Computes the gradient of the specified function together with the value."""
  if 'value_and_grad' in backend():
    return backend()['value_and_grad'](*args, **kwargs)
  grad_fn = grad(*args, **kwargs)
  fn = args[0]
  has_aux = False
  if has_aux in kwargs:
    has_aux = kwargs['has_aux']
  if not has_aux:
    def val_and_grad(*fn_args, **fn_kwargs):
      return fn(*fn_args, **fn_kwargs), grad_fn(*fn_args, **fn_kwargs)
    return val_and_grad
  def val_and_grad_aux(*fn_args, **fn_kwargs):
    g, aux = grad_fn(*fn_args, **fn_kwargs)
    res, _ = fn(*fn_args, **fn_kwargs)
    return (res, aux), g
  return val_and_grad_aux


def vjp(*args, **kwargs):
  """Computes the vector-Jacobian product for the specified function."""
  return backend()['vjp'](*args, **kwargs)


def custom_grad(*args, **kwargs):
  """Set a custom gradient computation (override the default) for a function."""
  return backend()['custom_grad'](*args, **kwargs)


def pmap(*args, **kwargs):
  """Parallel-map to apply a function on multiple accelerators in parallel."""
  return backend()['pmap'](*args, **kwargs)


def psum(*args, **kwargs):
  """Parallel-sum to use within a pmap'd function for aggregation."""
  return backend()['psum'](*args, **kwargs)


def abstract_eval(*args, **kwargs):
  """Evaluates function just on signatures of parameters, return signatures."""
  return backend()['abstract_eval'](*args, **kwargs)


def dataset_as_numpy(*args, **kwargs):
  """Convert a tf.data.Dataset to a stream of numpy arrays."""
  return backend()['dataset_as_numpy'](*args, **kwargs)


def device_count(*args, **kwargs):
  """Return the number of accelerators (GPUs or TPUs) available."""
  return backend()['device_count'](*args, **kwargs)


# Backend selection functions.


override_backend_name = None


@gin.configurable()
def backend(name='jax'):
  """Return the backend used to provide fastmath ops ('tf' or 'jax')."""
  name = name if not override_backend_name else override_backend_name
  if name == 'numpy':
    return NUMPY_BACKEND
  elif name == 'tf':
    return TF_BACKEND
  return JAX_BACKEND


@contextlib.contextmanager
def use_backend(name):
  """Call fastmath functions with a specified backend."""
  global override_backend_name
  prev_name = override_backend_name
  override_backend_name = name
  # Run the decorated function in try-finally in case it throws, e.g. for tests.
  try:
    yield
  finally:
    override_backend_name = prev_name


def backend_name():
  """Returns the name of the backend curently in use ('tf' or 'jax')."""
  return backend()['name']

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	"""Trax accelerated math operations for fast computing on GPUs and TPUs.	1✔
18
19	Import these operations directly from fastmath and import fastmath.numpy as np::
20
21	from trax import fastmath
22	from trax.fastmath import numpy as np
23
24	x = np.array([1.0, 2.0]) # Use like numpy.
25	y = np.exp(x) # Common numpy ops are available and accelerated.
26	z = fastmath.logsumexp(y) # Special operations available from fastmath.
27
28	Trax uses either TensorFlow 2 or JAX as backend for accelerating operations.
29	You can select which one to use (e.g., for debugging) with `use_backend`.
30	"""
31
32	import contextlib	1✔
33
34	import gin	1✔
35	from trax.fastmath.jax import JAX_BACKEND	1✔
36	from trax.fastmath.numpy import NUMPY_BACKEND	1✔
37	from trax.fastmath.tf import TF_BACKEND	1✔
38
39	# For numpy and random modules, we need to call "backend()" lazily, only when
40	# the function is called -- so that it can be set by gin configs.
41	# (Otherwise, backend() is called on import before gin-config is parsed.)
42	# To do that, we make objects to encapsulated these modules.
43
44
45	# A class that just forwards attribute accesses to backend's numpy object.
46	class NumpyBackend(object):	1✔
47	"""Numpy functions accelerated to run on GPUs and TPUs. Use like numpy."""
48
49	def __getattr__(self, attr):	1✔
50	return getattr(backend()['np'], attr)	1✔
51
52	numpy = NumpyBackend()	1✔
53
54
55	class RandomBackend(object):	1✔
56	"""Backend providing random functions."""
57
58	def get_prng(self, seed):	1✔
59	return backend()['random_get_prng'](seed)	1✔
60
61	def split(self, prng, num=2):	1✔
62	return backend()['random_split'](prng, num)	1✔
63
64	def uniform(self, args, *kwargs):	1✔
65	return backend()['random_uniform'](args, *kwargs)	1✔
66
67	def randint(self, args, *kwargs):	1✔
68	return backend()['random_randint'](args, *kwargs)	×
69
70	def normal(self, args, *kwargs):	1✔
71	return backend()['random_normal'](args, *kwargs)	1✔
72
73	def bernoulli(self, args, *kwargs):	1✔
74	return backend()['random_bernoulli'](args, *kwargs)	1✔
75
76
77	random = RandomBackend()	1✔
78
79
80	def logsumexp(args, *kwargs):	1✔
81	"""Computes the log of the sum of exponentials of input elements."""
82	return backend()['logsumexp'](args, *kwargs)	1✔
83
84
85	def expit(args, *kwargs):	1✔
86	"""Computes the expit (sigmoid) function."""
87	return backend()['expit'](args, *kwargs)	×
88
89
90	def sigmoid(args, *kwargs):	1✔
91	"""Computes the sigmoid (expit) function."""
92	return backend()['expit'](args, *kwargs)	×
93
94
95	def erf(args, *kwargs):	1✔
96	"""Computes the erf function."""
97	return backend()['erf'](args, *kwargs)	×
98
99
100	def conv(args, *kwargs):	1✔
101	"""Computes a generalized convolution."""
102	return backend()['conv'](args, *kwargs)	×
103
104
105	def avg_pool(args, *kwargs):	1✔
106	"""Average pooling."""
107	return backend()['avg_pool'](args, *kwargs)	×
108
109
110	def max_pool(args, *kwargs):	1✔
111	"""Max pooling."""
112	return backend()['max_pool'](args, *kwargs)	×
113
114
115	def sum_pool(args, *kwargs):	1✔
116	"""Sum pooling."""
117	return backend()['sum_pool'](args, *kwargs)	×
118
119
120	def sort_key_val(args, *kwargs):	1✔
121	"""Sorts keys along dimension and applies same permutation to values."""
122	return backend()['sort_key_val'](args, *kwargs)	×
123
124
125	def scan(args, *kwargs):	1✔
126	"""Scan to make recurrent functions run faster on accelerators."""
127	return backend()['scan'](args, *kwargs)	×
128
129
130	def remat(args, *kwargs):	1✔
131	"""Recompute everything in the backward pass to same memory."""
132	return backend()['remat'](args, *kwargs)	×
133
134
135	def cond(args, *kwargs):	1✔
136	"""Conditional computation to run on accelerators."""
137	return backend()['cond'](args, *kwargs)	×
138
139
140	def lt(args, *kwargs):	1✔
141	"""Less-than function for backends that do not override <."""
142	return backend()['lt'](args, *kwargs)	×
143
144
145	def stop_gradient(args, *kwargs):	1✔
146	"""Identity on the forward pass but 0 (no gradient) on the backward pass."""
147	return backend()['stop_gradient'](args, *kwargs)	×
148
149
150	_disable_jit = False	1✔
151
152
153	def jit(args, *kwargs):	1✔
154	"""Just-In-Time compiles the given function for use on accelerators."""
155	global _disable_jit
156	if _disable_jit:	×
157	return args[0] # jit(f, **unused_now_jit_kwargs) = f	×
158	return backend()['jit'](args, *kwargs)	×
159
160
161	def disable_jit():	1✔
162	"""Disables JIT-compilation; helpful for debugging."""
163	global _disable_jit
164	_disable_jit = True	×
165
166
167	def grad(args, *kwargs):	1✔
168	"""Computes the gradient of the specified function (returns a function)."""
169	return backend()['grad'](args, *kwargs)	×
170
171
172	def value_and_grad(args, *kwargs):	1✔
173	"""Computes the gradient of the specified function together with the value."""
174	if 'value_and_grad' in backend():	×
175	return backend()['value_and_grad'](args, *kwargs)	×
176	grad_fn = grad(args, *kwargs)	×
177	fn = args[0]	×
178	has_aux = False	×
179	if has_aux in kwargs:	×
180	has_aux = kwargs['has_aux']	×
181	if not has_aux:	×
182	def val_and_grad(fn_args, *fn_kwargs):	×
183	return fn(fn_args, fn_kwargs), grad_fn(fn_args, **fn_kwargs)	×
184	return val_and_grad	×
185	def val_and_grad_aux(fn_args, *fn_kwargs):	×
186	g, aux = grad_fn(fn_args, *fn_kwargs)	×
187	res, _ = fn(fn_args, *fn_kwargs)	×
188	return (res, aux), g	×
189	return val_and_grad_aux	×
190
191
192	def vjp(args, *kwargs):	1✔
193	"""Computes the vector-Jacobian product for the specified function."""
194	return backend()['vjp'](args, *kwargs)	×
195
196
197	def custom_grad(args, *kwargs):	1✔
198	"""Set a custom gradient computation (override the default) for a function."""
199	return backend()['custom_grad'](args, *kwargs)	×
200
201
202	def pmap(args, *kwargs):	1✔
203	"""Parallel-map to apply a function on multiple accelerators in parallel."""
204	return backend()['pmap'](args, *kwargs)	×
205
206
207	def psum(args, *kwargs):	1✔
208	"""Parallel-sum to use within a pmap'd function for aggregation."""
209	return backend()['psum'](args, *kwargs)	×
210
211
212	def abstract_eval(args, *kwargs):	1✔
213	"""Evaluates function just on signatures of parameters, return signatures."""
214	return backend()['abstract_eval'](args, *kwargs)	1✔
215
216
217	def dataset_as_numpy(args, *kwargs):	1✔
218	"""Convert a tf.data.Dataset to a stream of numpy arrays."""
219	return backend()['dataset_as_numpy'](args, *kwargs)	×
220
221
222	def device_count(args, *kwargs):	1✔
223	"""Return the number of accelerators (GPUs or TPUs) available."""
224	return backend()['device_count'](args, *kwargs)	×
225
226
227	# Backend selection functions.
228
229
230	override_backend_name = None	1✔
231
232
233	@gin.configurable()	1✔
234	def backend(name='jax'):	1✔
235	"""Return the backend used to provide fastmath ops ('tf' or 'jax')."""
236	name = name if not override_backend_name else override_backend_name	1✔
237	if name == 'numpy':	1✔
238	return NUMPY_BACKEND	×
239	elif name == 'tf':	1✔
240	return TF_BACKEND	1✔
241	return JAX_BACKEND	1✔
242
243
244	@contextlib.contextmanager	1✔
245	def use_backend(name):
246	"""Call fastmath functions with a specified backend."""
247	global override_backend_name
248	prev_name = override_backend_name	1✔
249	override_backend_name = name	1✔
250	# Run the decorated function in try-finally in case it throws, e.g. for tests.
251	try:	1✔
252	yield	1✔
253	finally:
254	override_backend_name = prev_name	1✔
255
256
257	def backend_name():	1✔
258	"""Returns the name of the backend curently in use ('tf' or 'jax')."""
259	return backend()['name']	1✔

google / trax / 2733

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