wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/ops/numpy_ops/tests/BUILD
+284
View File
@@ -0,0 +1,284 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
licenses(["notice"])
py_library(
name = "config",
srcs = ["config.py"],
strict_deps = True,
deps = [
"@absl_py//absl:app",
"@absl_py//absl/flags",
],
)
py_library(
name = "test_util",
srcs = ["test_util.py"],
strict_deps = True,
deps = [
":config",
":extensions",
"//tensorflow:tensorflow_py",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:gradient_checker_v2",
"//tensorflow/python/ops/numpy_ops:np_array_ops",
"//tensorflow/python/ops/numpy_ops:np_utils",
"//tensorflow/python/ops/numpy_ops:numpy",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:numpy_compat",
"//third_party/py/numpy",
"@absl_py//absl/testing:absltest",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "np_wrapper",
srcs = ["np_wrapper.py"],
strict_deps = True,
visibility = [
"//visibility:public",
],
deps = [
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops/numpy_ops:np_array_ops",
"//tensorflow/python/ops/numpy_ops:np_arrays",
"//tensorflow/python/ops/numpy_ops:np_config",
"//tensorflow/python/ops/numpy_ops:np_dtypes",
"//tensorflow/python/ops/numpy_ops:np_math_ops",
"//tensorflow/python/ops/numpy_ops:np_random",
"//tensorflow/python/ops/numpy_ops:np_utils",
"//tensorflow/python/ops/numpy_ops:numpy",
"//third_party/py/numpy",
],
)
py_library(
name = "extensions",
srcs = ["extensions.py"],
strict_deps = True,
deps = [
":np_wrapper",
"//tensorflow:tensorflow_py",
"//tensorflow/python/compiler/xla",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager/polymorphic_function",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:device_spec",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_conversion_registry",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:bitwise_ops_gen",
"//tensorflow/python/ops:clip_ops",
"//tensorflow/python/ops:collective_ops_gen",
"//tensorflow/python/ops:control_flow_assert",
"//tensorflow/python/ops:custom_gradient",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:special_math_ops",
"//tensorflow/python/ops:stateless_random_ops",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/ops/numpy_ops:numpy",
"//tensorflow/python/ops/parallel_for:control_flow_ops",
"//tensorflow/python/tpu:tpu_py",
"//tensorflow/python/tpu/ops",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
"@six_archive//:six",
],
)
# copybara:uncomment_begin(google-only)
# py_test(
# name = "extensions_test",
# srcs = ["extensions_test.py"],
# strict_deps = True,
# tags = [
# "gpu",
# "no_pip",
# "notap", # b/294137902
# "requires-gpu-nvidia",
# ],
# deps = [
# ":extensions",
# ":np_wrapper",
# "//learning/brain/research/jax:gpu_support",
# # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
# "//third_party/py/jax",
# "//third_party/py/numpy",
# "//tensorflow:tensorflow_py",
# "//tensorflow/python/distribute/cluster_resolver/tpu:tpu_cluster_resolver_py",
# "//tensorflow/python/eager:backprop",
# "//tensorflow/python/eager:context",
# "//tensorflow/python/eager/polymorphic_function",
# "//tensorflow/python/framework:config",
# "//tensorflow/python/framework:constant_op",
# "//tensorflow/python/framework:dtypes",
# "//tensorflow/python/framework:ops",
# "//tensorflow/python/framework:tensor",
# "//tensorflow/python/framework:tensor_shape",
# "//tensorflow/python/ops:array_ops",
# "//tensorflow/python/ops:gradient_checker_v2",
# "//tensorflow/python/ops:math_ops",
# "//tensorflow/python/ops:nn_ops",
# "//tensorflow/python/ops:random_ops",
# "//tensorflow/python/ops:stateful_random_ops",
# "//tensorflow/python/ops/numpy_ops:numpy",
# "//tensorflow/python/platform:client_testlib",
# "//tensorflow/python/util:nest",
# "@absl_py//absl/flags",
# "@absl_py//absl/testing:parameterized",
# ],
# )
#
# py_test(
# name = "extensions_test_tpu",
# srcs = ["extensions_test.py"],
# args = [
# "--jax_allow_unused_tpus",
# "--requires_tpu",
# ],
# main = "extensions_test.py",
# strict_deps = True,
# tags = [
# "no_pip",
# "requires-tpu",
# ],
# deps = [
# ":extensions",
# ":np_wrapper",
# "//learning/brain/google/xla",
# # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
# "//third_party/py/jax",
# "//third_party/py/numpy",
# "//tensorflow:tensorflow_py",
# "//tensorflow/python/distribute/cluster_resolver/tpu:tpu_cluster_resolver_py",
# "//tensorflow/python/eager:backprop",
# "//tensorflow/python/eager:context",
# "//tensorflow/python/eager/polymorphic_function",
# "//tensorflow/python/framework:config",
# "//tensorflow/python/framework:constant_op",
# "//tensorflow/python/framework:dtypes",
# "//tensorflow/python/framework:ops",
# "//tensorflow/python/framework:tensor",
# "//tensorflow/python/framework:tensor_shape",
# "//tensorflow/python/ops:array_ops",
# "//tensorflow/python/ops:gradient_checker_v2",
# "//tensorflow/python/ops:math_ops",
# "//tensorflow/python/ops:nn_ops",
# "//tensorflow/python/ops:random_ops",
# "//tensorflow/python/ops:stateful_random_ops",
# "//tensorflow/python/ops/numpy_ops:numpy",
# "//tensorflow/python/platform:client_testlib",
# "//tensorflow/python/util:nest",
# "@absl_py//absl/flags",
# "@absl_py//absl/testing:parameterized",
# ],
# )
# copybara:uncomment_end
py_test(
name = "np_test",
timeout = "long",
srcs = ["np_test.py"],
args = [
"--num_generated_cases=90",
"--enable_x64", # Needed to enable dtype check
],
shard_count = 20,
strict_deps = True,
tags = [
"gpu",
"no_pip",
"requires-gpu-nvidia",
],
deps = [
":config",
":extensions",
":np_wrapper",
":test_util",
# copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
"//third_party/py/numpy",
"@six_archive//:six",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops/numpy_ops:np_config",
"//tensorflow/python/ops/numpy_ops:numpy",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:numpy_compat",
"@absl_py//absl/testing:absltest",
"@absl_py//absl/testing:parameterized",
],
)
py_test(
name = "np_indexing_test",
srcs = ["np_indexing_test.py"],
args = [
"--num_generated_cases=90",
"--enable_x64", # Needed to enable dtype check
],
shard_count = 10,
strict_deps = True,
# TODO(b/164245103): Re-enable GPU once tf.tensor_strided_slice_update's segfault is fixed.
tags = [
"no_pip",
# "gpu",
# "requires-gpu-nvidia",
],
deps = [
":config",
":extensions",
":np_wrapper",
":test_util",
# copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
"//third_party/py/numpy",
"//tensorflow/python/framework:config",
"//tensorflow/python/ops/numpy_ops:numpy",
"//tensorflow/python/util:nest",
"@absl_py//absl/testing:absltest",
"@absl_py//absl/testing:parameterized",
],
)
py_test(
name = "np_einsum_test",
srcs = ["np_einsum_test.py"],
args = [
"--num_generated_cases=90",
"--enable_x64", # Needed to enable dtype check
],
shard_count = 20,
strict_deps = True,
tags = [
"gpu",
"no_pip",
"requires-gpu-nvidia",
],
deps = [
":config",
":np_wrapper",
":test_util",
# copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
"//third_party/py/numpy",
"//tensorflow/python/ops/numpy_ops:np_config",
"//tensorflow/python/ops/numpy_ops:numpy",
"@absl_py//absl/testing:absltest",
"@absl_py//absl/testing:parameterized",
],
)
tensorflow/python/ops/numpy_ops/tests/config.py
+141
View File
@@ -0,0 +1,141 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Test configurations."""
import os
import sys
def bool_env(varname: str, default: bool) -> bool:
"""Read an environment variable and interpret it as a boolean.
True values are (case insensitive): 'y', 'yes', 't', 'true', 'on', and '1';
false values are 'n', 'no', 'f', 'false', 'off', and '0'.
Args:
varname: the name of the variable
default: the default boolean value
Raises: ValueError if the environment variable is anything else.
"""
val = os.getenv(varname, str(default))
val = val.lower()
if val in ('y', 'yes', 't', 'true', 'on', '1'):
return True
elif val in ('n', 'no', 'f', 'false', 'off', '0'):
return False
else:
raise ValueError(
'invalid truth value %r for environment %r' % (val, varname)
)
class Config(object):
def __init__(self):
self.values = {}
self.meta = {}
self.FLAGS = NameSpace(self.read)
self.use_absl = False
def update(self, name, val):
if self.use_absl:
setattr(self.absl_flags.FLAGS, name, val)
else:
self.check_exists(name)
if name not in self.values:
raise Exception("Unrecognized config option: {}".format(name))
self.values[name] = val
def read(self, name):
if self.use_absl:
return getattr(self.absl_flags.FLAGS, name)
else:
self.check_exists(name)
return self.values[name]
def add_option(self, name, default, opt_type, meta_args, meta_kwargs):
if name in self.values:
raise Exception("Config option {} already defined".format(name))
self.values[name] = default
self.meta[name] = (opt_type, meta_args, meta_kwargs)
def check_exists(self, name):
if name not in self.values:
raise Exception("Unrecognized config option: {}".format(name))
def DEFINE_bool(self, name, default, *args, **kwargs):
self.add_option(name, default, bool, args, kwargs)
def DEFINE_integer(self, name, default, *args, **kwargs):
self.add_option(name, default, int, args, kwargs)
def DEFINE_string(self, name, default, *args, **kwargs):
self.add_option(name, default, str, args, kwargs)
def DEFINE_enum(self, name, default, *args, **kwargs):
self.add_option(name, default, 'enum', args, kwargs)
def config_with_absl(self):
# Run this before calling `app.run(main)` etc
import absl.flags as absl_FLAGS
from absl import app, flags as absl_flags
self.use_absl = True
self.absl_flags = absl_flags
absl_defs = { bool: absl_flags.DEFINE_bool,
int: absl_flags.DEFINE_integer,
str: absl_flags.DEFINE_string,
'enum': absl_flags.DEFINE_enum }
for name, val in self.values.items():
flag_type, meta_args, meta_kwargs = self.meta[name]
absl_defs[flag_type](name, val, *meta_args, **meta_kwargs)
app.call_after_init(lambda: self.complete_absl_config(absl_flags))
def complete_absl_config(self, absl_flags):
for name, _ in self.values.items():
self.update(name, getattr(absl_flags.FLAGS, name))
def parse_flags_with_absl(self):
global already_configured_with_absl
if not already_configured_with_absl:
import absl.flags
self.config_with_absl()
absl.flags.FLAGS(sys.argv, known_only=True)
self.complete_absl_config(absl.flags)
already_configured_with_absl = True
class NameSpace(object):
def __init__(self, getter):
self._getter = getter
def __getattr__(self, name):
return self._getter(name)
config = Config()
flags = config
FLAGS = flags.FLAGS
already_configured_with_absl = False
flags.DEFINE_bool(
'jax_enable_checks',
bool_env('JAX_ENABLE_CHECKS', False),
help='Turn on invariant checking (core.skip_checks = False)')
flags.DEFINE_bool('tf_numpy_additional_tests', True,
'Run tests added specifically for TF numpy')
tensorflow/python/ops/numpy_ops/tests/extensions.py
+2094
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/numpy_ops/tests/np_einsum_test.py
+355
View File
@@ -0,0 +1,355 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
from collections import defaultdict # pylint: disable=g-importing-member
import itertools
from absl.testing import absltest
from absl.testing import parameterized
import numpy as np
from tensorflow.python.ops.numpy_ops.tests.config import config
import tensorflow.python.ops.numpy_ops.tests.np_wrapper as tnp
import tensorflow.python.ops.numpy_ops.tests.test_util as tntu
config.parse_flags_with_absl()
class EinsumTest(tntu.TestCase):
def _check(self, s, *ops):
a = np.einsum(s, *ops)
b = tnp.einsum(s, *ops)
self.assertAllClose(a, b, check_dtypes=True, atol=1e-4, rtol=1e-4)
def test_three_operands_1(self):
r = self.rng()
x = r.randn(3)
y = r.randn(4)
z = r.randn(5)
s = 'i,j,k->ijk'
self._check(s, x, y, z)
def test_three_operands_2(self):
r = self.rng()
x = r.randn(3)
y = r.randn(4)
z = r.randn(5)
s = 'i,j,k->ijk'
self._check(s, x, y, z)
def test_two_operands_1(self):
r = self.rng()
x = r.randn(3, 4)
y = r.randn(4)
s = 'ij,j->i'
self._check(s, x, y)
def test_two_operands_2(self):
r = self.rng()
x = r.randn(3, 4, 5)
y = r.randn(4)
s = 'ijk,j->i'
self._check(s, x, y)
def test_two_operands_3(self):
r = self.rng()
x = r.randn(3, 4, 3)
y = r.randn(3)
s = 'iji,i->j'
self._check(s, x, y)
def test_two_operands_4(self):
r = self.rng()
x = r.randn(3, 4)
y = r.randn(3, 4)
s = 'ij,ij->'
self._check(s, x, y)
def test_two_operands_5(self):
r = self.rng()
x = r.randn(10, 2, 3)
y = r.randn(3, 4)
s = 'nij,jk->nik'
self._check(s, x, y)
def test_two_operands_6(self):
# based on https://github.com/google/jax/issues/37#issuecomment-448572187
r = self.rng()
x = r.randn(2, 1)
y = r.randn(2, 3, 4)
s = 'sa,shb->shab'
self._check(s, x, y)
def test_one_operand_1(self):
r = self.rng()
x = r.randn(3, 4, 5)
s = 'ijk->j'
self._check(s, x)
def test_one_operand_2(self):
r = self.rng()
x = r.randn(3, 4, 5)
s = 'ijk->kij'
self._check(s, x)
def test_one_operand_3(self):
r = self.rng()
x = r.randn(3, 4, 5)
s = 'ijk->ki'
self._check(s, x)
def test_one_operand_4(self):
r = self.rng()
x = r.randn(3, 4, 5)
s = 'ijk->ki'
self._check(s, x)
def test_one_operand_5(self):
r = self.rng()
x = r.randn(2, 3, 4, 5)
s = '...ijk->...ki'
self._check(s, x)
def test_one_operand_6(self):
r = self.rng()
x = r.randn(3, 4, 5)
s = '...ijk->ki'
self._check(s, x)
def test_one_operand_7(self):
r = self.rng()
x = r.randn(3, 3)
s = 'ii->'
self._check(s, x)
def test_one_operand_8(self):
r = self.rng()
x = r.randn(3, 3)
s = 'ij->'
self._check(s, x)
def test_one_operand_9(self):
r = self.rng()
x = r.randn(3, 3, 3)
s = 'iii->'
self._check(s, x)
def test_one_operand_10(self):
r = self.rng()
x = r.randn(3, 3)
s = 'ii->i'
self._check(s, x)
def test_one_operand_11(self):
r = self.rng()
x = r.randn(3, 3, 4)
s = 'iij->i'
self._check(s, x)
def test_one_operand_12(self):
r = self.rng()
x = r.randn(3, 3, 3)
s = 'iii->i'
self._check(s, x)
def test_one_operand_13(self):
r = self.rng()
x = r.randn(3, 3, 5, 4, 4)
s = 'iijkk->i'
self._check(s, x)
def test_one_operand_14(self):
r = self.rng()
x = r.randn(3, 3, 5, 4, 4)
s = 'iijkk->ik'
self._check(s, x)
def test_one_operand_15(self):
r = self.rng()
x = r.randn(3, 3, 5, 4, 4)
s = 'iijkl->il'
self._check(s, x)
def test_one_operand_16(self):
r = self.rng()
x = r.randn(3, 3)
s = 'ij->ij'
self._check(s, x)
def test_tf_unsupported_1(self):
# from https://www.tensorflow.org/api_docs/python/tf/einsum
r = self.rng()
x = r.randn(2, 3, 5, 1)
y = r.randn(3, 4, 5, 1)
s = 'ij...,jk...->ik...'
self._check(s, x, y)
def test_tf_unsupported_2(self):
# from https://www.tensorflow.org/api_docs/python/tf/einsum
r = self.rng()
x = r.randn(2, 3, 3)
y = r.randn(4)
s = 'ijj,k->ik'
self._check(s, x, y)
def test_tf_unsupported_3(self):
# from https://www.tensorflow.org/api_docs/python/tf/einsum
r = self.rng()
x = r.randn(2, 3)
y = r.randn(2, 3)
z = r.randn(3, 4)
s = 'ij,ij,jk->ik'
self._check(s, x, y, z)
# these tests are based on https://github.com/dask/dask/pull/3412/files
@parameterized.named_parameters(
{'testcase_name': '_{}_dtype={}'.format(einstr, dtype.__name__), # pylint: disable=g-complex-comprehension
'einstr': einstr, 'dtype': dtype}
for einstr in [
'abc,bad->abcd',
'abcdef,bcdfg->abcdeg',
'ea,fb,abcd,gc,hd->efgh',
'ab,b',
'aa',
'a,a->',
'a,a->a',
'a,a',
'a,b',
'a,b,c',
'a',
'ba,b',
'ba,b->',
'defab,fedbc->defac',
'ab...,bc...->ac...',
'a...a',
'abc...->cba...',
'...ab->...a',
'a...a->a...',
# Following 2 from # https://stackoverflow.com/a/19203475/1611416
'...abc,...abcd->...d',
'ab...,b->ab...',
# https://github.com/dask/dask/pull/3412#discussion_r182413444
'aa->a',
'ab,ab,c->c',
'aab,bc->ac',
'aab,bcc->ac',
'fdf,cdd,ccd,afe->ae',
'fff,fae,bef,def->abd',
]
# TODO(wangpeng): Add tnp.bool_ to dtype list
for dtype in [tnp.float32, tnp.int32, tnp.complex64])
def test_from_dask(self, einstr, dtype):
r = tntu.rand_default()
if '->' in einstr:
input_str, _ = einstr.split('->')
else:
input_str = einstr
input_names = input_str.split(',')
dims = itertools.cycle([2, 3, 4])
shapes = defaultdict(lambda: next(dims))
input_shapes = [tuple(shapes[c] for c in names.replace('...', '01'))
for names in input_names]
operands = [r(shape, dtype) for shape in input_shapes]
self._check(einstr, *operands)
def test_ordered_front_batch_dim_case(self):
x = np.ones((1, 8, 20, 4))
y = np.ones((1, 8, 20, 4))
s = 'ijkl,ijml->ijkm'
self._check(s, x, y)
# pylint: disable=invalid-name
def test_einsum_path(self):
# just check examples from np.einsum_path docstring
a = self.rng().rand(2, 2)
b = self.rng().rand(2, 5)
c = self.rng().rand(5, 2)
path_info = np.einsum_path('ij,jk,kl->il', a, b, c, optimize='greedy')
self.assertEqual(str(path_info[0]), "['einsum_path', (1, 2), (0, 1)]")
self.assertEqual(path_info[1].split('\n')[0],
' Complete contraction: ij,jk,kl->il')
# check this doesn't crash
I = self.rng().rand(10, 10, 10, 10)
C = self.rng().rand(10, 10)
np.einsum_path('ea,fb,abcd,gc,hd->efgh', C, C, I, C, C, optimize='greedy')
@tntu.disable
def test_einsum_kpmurphy_example(self):
# code from an email with @murphyk
N = 2
C = 3
D = 4
K = 5
T = 6
r = self.rng()
S = r.randn(N, T, K)
W = r.randn(K, D)
V = r.randn(D, C)
L = np.zeros((N, C))
for n in range(N):
for c in range(C):
s = 0
for d in range(D):
for k in range(K):
for t in range(T):
s += S[n, t, k] * W[k, d] * V[d, c]
L[n, c] = s
path = tnp.einsum_path('ntk,kd,dc->nc', S, W, V, optimize='optimal')[0]
rtol = 1e-2 if tntu.device_under_test() == 'tpu' else None
self.assertAllClose(L, tnp.einsum('ntk,kd,dc->nc', S, W, V, optimize=path),
check_dtypes=False, rtol=rtol)
# pylint: enable=invalid-name
@tntu.disable
def test_contraction_broadcasting(self):
r = self.rng()
x = r.randn(3, 4, 5)
y = r.randn(3, 1, 6)
s = 'cij,cjk->cik'
self._check(s, x, y)
@tntu.disable
def test_batch_broadcasting(self):
r = self.rng()
x = r.randn(1, 4, 5)
y = r.randn(3, 5, 6)
s = 'cij,cjk->cik'
self._check(s, x, y)
@tntu.disable
def test_batch_and_contraction_broadcasting(self):
r = self.rng()
x = r.randn(1, 4, 5)
y = r.randn(3, 1, 6)
s = 'cij,cjk->cik'
self._check(s, x, y)
@tntu.disable
def test_broadcasting_issue_2189(self):
r = self.rng()
x = r.randn(2, 1, 3, 3)
y = r.randn(2, 4, 3)
s = '...ij,...j'
self._check(s, x, y)
if __name__ == '__main__':
absltest.main()
tensorflow/python/ops/numpy_ops/tests/np_indexing_test.py
+988
View File
@@ -0,0 +1,988 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
import collections
import enum
from functools import partial
import itertools
import unittest
from absl.testing import absltest
from absl.testing import parameterized
import numpy as onp
from tensorflow.python.framework import config as tf_config
from tensorflow.python.ops.numpy_ops.tests.config import config
import tensorflow.python.ops.numpy_ops.tests.extensions as nje
import tensorflow.python.ops.numpy_ops.tests.np_wrapper as tnp
import tensorflow.python.ops.numpy_ops.tests.test_util as jtu
from tensorflow.python.util import nest
config.parse_flags_with_absl()
# We disable the whitespace continuation check in this file because otherwise it
# makes the test name formatting unwieldy.
# pylint: disable=bad-continuation
# We also disable undefined-variable till we start enabling tests.
# pylint: disable=undefined-variable
def subvals(lst, replace):
lst = list(lst)
for i, v in replace:
lst[i] = v
return tuple(lst)
float_dtypes = [onp.float32, onp.float64]
int_dtypes = [onp.int32, onp.int64]
bool_types = [onp.bool_]
default_dtypes = float_dtypes + int_dtypes
all_dtypes = float_dtypes + int_dtypes + bool_types
IndexSpec = collections.namedtuple("IndexTest", ["shape", "indexer"])
suppress_deprecated_indexing_warnings = partial(
jtu.ignore_warning, category=FutureWarning,
message='Using a non-tuple sequence.*')
STATIC_INDEXING_TESTS = [
("OneIntIndex", [
IndexSpec(shape=(3,), indexer=1),
IndexSpec(shape=(3, 3), indexer=0),
IndexSpec(shape=(3, 4, 5), indexer=2),
IndexSpec(shape=(3,), indexer=-1),
IndexSpec(shape=(3,), indexer=-2),
]),
("TwoIntIndices", [
IndexSpec(shape=(3, 3), indexer=(2, 1)),
IndexSpec(shape=(3, 4, 5), indexer=(1, 2)),
IndexSpec(shape=(3, 4, 5), indexer=(-1, 2)),
]),
("ThreeIntIndices", [IndexSpec((3, 4, 5), indexer=(1, 2, 3))]),
("OneSliceIndex", [
IndexSpec(shape=(10,), indexer=slice(1, 3)),
IndexSpec(shape=(10,), indexer=slice(1, -1)),
IndexSpec(shape=(10,), indexer=slice(None, -1)),
IndexSpec(shape=(10,), indexer=slice(None, None, None)),
IndexSpec(shape=(10, 8), indexer=slice(1, 3)),
IndexSpec(shape=(10, 8), indexer=slice(1, None)),
IndexSpec(shape=(10, 8), indexer=slice(None, 3)),
IndexSpec(shape=(10, 8), indexer=slice(-3, None)),
]),
("OneSliceIndexNegativeStride", [
IndexSpec(shape=(10,), indexer=slice(3, 1, -1)),
IndexSpec(shape=(10,), indexer=slice(1, 8, -1)), # empty result
IndexSpec(shape=(10,), indexer=slice(None, 1, -2)),
IndexSpec(shape=(10,), indexer=slice(None, None, -1)),
IndexSpec(shape=(10, 8), indexer=slice(3, 1, -1)),
IndexSpec(shape=(10, 8), indexer=slice(0, 8, -1)), # empty result
IndexSpec(shape=(10, 8), indexer=slice(None, None, -1)),
]),
("OneSliceIndexNonUnitStride", [
IndexSpec(shape=(10,), indexer=slice(0, 8, 2)),
IndexSpec(shape=(10,), indexer=slice(0, 8, 3)),
IndexSpec(shape=(10,), indexer=slice(1, 3, 2)),
IndexSpec(shape=(10,), indexer=slice(1, None, 2)),
IndexSpec(shape=(10,), indexer=slice(None, 1, -2)),
IndexSpec(shape=(10, 8), indexer=slice(1, 8, 3)),
IndexSpec(shape=(10, 8), indexer=slice(None, None, 2)),
IndexSpec(shape=(10, 8), indexer=slice(None, 1, -2)),
IndexSpec(shape=(10, 8), indexer=slice(None, None, -2)),
]),
("TwoSliceIndices", [
IndexSpec(shape=(10, 8), indexer=(slice(1, 3), slice(0, 2))),
IndexSpec(shape=(10, 8), indexer=(slice(1, None), slice(None, 2))),
IndexSpec(
shape=(10, 8), indexer=(slice(None, None, -1), slice(None, 2))),
IndexSpec(shape=(10, 8, 3), indexer=(slice(1, 3), slice(0, 2))),
IndexSpec(shape=(10, 8, 3), indexer=(slice(1, 3), slice(0, None))),
IndexSpec(shape=(10, 8, 3), indexer=(slice(1, None), slice(0, 2))),
]),
("OneColonIndex", [
IndexSpec(shape=(3,), indexer=slice(None)),
IndexSpec(shape=(3, 4), indexer=slice(None)),
]),
("MultipleColonIndices", [
IndexSpec(shape=(3, 4), indexer=(slice(None), slice(None))),
IndexSpec(shape=(3, 4, 5), indexer=(slice(None), slice(None))),
]),
("MixedSliceIndices", [
IndexSpec(shape=(10, 4), indexer=(slice(None), slice(0, 2))),
IndexSpec(shape=(10, 4), indexer=(1, slice(None))),
]),
("EllipsisIndex", [
IndexSpec(shape=(3,), indexer=Ellipsis),
IndexSpec(shape=(3, 4), indexer=Ellipsis),
IndexSpec(shape=(3, 4, 5), indexer=(0, Ellipsis)),
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis, 2, 3)),
]),
("NoneIndex", [
IndexSpec(shape=(), indexer=None),
IndexSpec(shape=(), indexer=(None, None)),
IndexSpec(shape=(), indexer=(Ellipsis, None)),
IndexSpec(shape=(3,), indexer=None),
IndexSpec(shape=(3, 4), indexer=None),
IndexSpec(shape=(3, 4), indexer=(Ellipsis, None)),
IndexSpec(shape=(3, 4), indexer=(0, None, Ellipsis)),
IndexSpec(shape=(3, 4, 5), indexer=(1, None, Ellipsis)),
]),
("EmptyIndex", [
IndexSpec(shape=(), indexer=()),
IndexSpec(shape=(3,), indexer=()),
IndexSpec(shape=(3, 4), indexer=()),
]),
]
STATIC_INDEXING_GRAD_TESTS = [
("OneIntIndex", [
IndexSpec(shape=(3,), indexer=1),
IndexSpec(shape=(3, 3), indexer=0),
IndexSpec(shape=(3, 4, 5), indexer=2),
IndexSpec(shape=(3,), indexer=-1),
IndexSpec(shape=(3,), indexer=-2),
]),
("TwoIntIndices", [
IndexSpec(shape=(3, 3), indexer=(2, 1)),
IndexSpec(shape=(3, 4, 5), indexer=(1, 2)),
IndexSpec(shape=(3, 4, 5), indexer=(-1, 2)),
]),
("ThreeIntIndices", [IndexSpec((3, 4, 5), indexer=(1, 2, 3))]),
("OneSliceIndex", [
IndexSpec(shape=(5,), indexer=slice(1, 3)),
IndexSpec(shape=(5,), indexer=slice(1, -1)),
IndexSpec(shape=(5,), indexer=slice(None, -1)),
IndexSpec(shape=(5,), indexer=slice(None, None, None)),
IndexSpec(shape=(5, 4), indexer=slice(1, 3)),
IndexSpec(shape=(5, 4), indexer=slice(1, None)),
IndexSpec(shape=(5, 4), indexer=slice(None, 3)),
IndexSpec(shape=(5, 4), indexer=slice(-3, None)),
]),
("TwoSliceIndices", [
IndexSpec(shape=(5, 4), indexer=(slice(1, 3), slice(0, 2))),
IndexSpec(shape=(5, 4), indexer=(slice(1, None), slice(None, 2))),
IndexSpec(shape=(5, 4, 3), indexer=(slice(1, 3), slice(0, 2))),
IndexSpec(shape=(5, 4, 3), indexer=(slice(1, 3), slice(0, None))),
IndexSpec(shape=(5, 4, 3), indexer=(slice(1, None), slice(0, 2))),
]),
("OneColonIndex", [
IndexSpec(shape=(3,), indexer=slice(None)),
IndexSpec(shape=(3, 4), indexer=slice(None)),
]),
("MultipleColonIndices", [
IndexSpec(shape=(3, 4), indexer=(slice(None), slice(None))),
IndexSpec(shape=(3, 4, 5), indexer=(slice(None), slice(None))),
]),
("MixedSliceIndices", [
IndexSpec(shape=(5, 4), indexer=(slice(None), slice(0, 2))),
IndexSpec(shape=(5, 4), indexer=(1, slice(None))),
]),
("EllipsisIndex", [
IndexSpec(shape=(3,), indexer=Ellipsis),
IndexSpec(shape=(3, 4), indexer=Ellipsis),
IndexSpec(shape=(3, 4, 5), indexer=(0, Ellipsis)),
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis, 2, 3)),
]),
("NoneIndex", [
IndexSpec(shape=(), indexer=None),
IndexSpec(shape=(), indexer=(None, None)),
IndexSpec(shape=(), indexer=(Ellipsis, None)),
IndexSpec(shape=(3,), indexer=None),
IndexSpec(shape=(3, 4), indexer=None),
IndexSpec(shape=(3, 4), indexer=(Ellipsis, None)),
IndexSpec(shape=(3, 4), indexer=(0, None, Ellipsis)),
IndexSpec(shape=(3, 4, 5), indexer=(1, None, Ellipsis)),
]),
# TODO(mattjj): these fail for uninteresting dtype reasons
# ("EmptyIndex",
# [IndexSpec(shape=(), indexer=()),
# IndexSpec(shape=(3,), indexer=()),
# IndexSpec(shape=(3, 4), indexer=()),
# ]),
]
ADVANCED_INDEXING_TESTS = [
("One1DIntArrayIndex",
[IndexSpec(shape=(3,), indexer=onp.array([0, 1])),
IndexSpec(shape=(3, 3), indexer=onp.array([1, 2, 1])),
IndexSpec(shape=(3, 4, 5), indexer=onp.array([0, 2, 0, 1])),
IndexSpec(shape=(3,), indexer=onp.array([-1, 1])),
IndexSpec(shape=(3,), indexer=onp.array([-2, -1])),
IndexSpec(shape=(0,), indexer=onp.array([], dtype=onp.int32)),
]),
("One2DIntArrayIndex",
[IndexSpec(shape=(3,), indexer=onp.array([[0, 0]])),
IndexSpec(shape=(3, 3), indexer=onp.array([[1, 2, 1],
[0, 1, -1]])),
IndexSpec(shape=(3, 4, 5), indexer=onp.array([[0, 2, 0, 1],
[-1, -2, 1, 0]])),
]),
("Two1DIntArrayIndicesNoBroadcasting",
[IndexSpec(shape=(3, 3), indexer=(onp.array([0, 1]),
onp.array([1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2, 0, 1]),
onp.array([-1, 0, -1, 2]))),
]),
("Two1DIntArrayIndicesWithBroadcasting",
[IndexSpec(shape=(3, 3), indexer=(onp.array([[0, 1]]),
onp.array([1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([[0, 2, 0, 1]]),
onp.array([-1, 0, -1, 2]))),
]),
("TupleOfListsOfPythonInts",
[IndexSpec(shape=(3, 4, 5), indexer=([0, 1])),
IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]], [[2, 3, 0, 3]])),
]),
("TupleOfPythonIntsAndIntArrays",
[IndexSpec(shape=(3, 4, 5), indexer=(0, onp.array([0, 1]))),
IndexSpec(shape=(3, 4, 5), indexer=(0, 1,
onp.array([[2, 3, 0, 3]]))),
]),
("TupleOfListsOfPythonIntsAndIntArrays",
[IndexSpec(shape=(3, 4, 5), indexer=([0, 1], onp.array([0]))),
IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]],
onp.array([[2, 3, 0, 3]]))),
]),
]
ADVANCED_INDEXING_TESTS_NO_REPEATS = [
("One1DIntArrayIndex",
[IndexSpec(shape=(3,), indexer=onp.array([0, 1])),
IndexSpec(shape=(3, 3), indexer=onp.array([1, 2, 0])),
IndexSpec(shape=(3, 4, 5), indexer=onp.array([0, 2, 1])),
IndexSpec(shape=(3,), indexer=onp.array([-1, 1])),
IndexSpec(shape=(3,), indexer=onp.array([-2, -1])),
# Fails with a TF/XLA error.
# IndexSpec(shape=(0,), indexer=onp.array([], dtype=onp.int32)),
]),
("One2DIntArrayIndex",
[IndexSpec(shape=(3,), indexer=onp.array([[0, 1]])),
IndexSpec(shape=(6, 6), indexer=onp.array([[1, 2, 0],
[3, 4, -1]])),
]),
("Two1DIntArrayIndicesNoBroadcasting",
[IndexSpec(shape=(3, 3), indexer=(onp.array([0, 1]),
onp.array([1, 2]))),
IndexSpec(shape=(4, 5, 6), indexer=(onp.array([0, 2, 1, 3]),
onp.array([-1, 0, -2, 1]))),
]),
("Two1DIntArrayIndicesWithBroadcasting",
[IndexSpec(shape=(3, 3), indexer=(onp.array([[0, 1]]),
onp.array([1, 2]))),
IndexSpec(shape=(4, 5, 6), indexer=(onp.array([[0, 2, -1, 1]]),
onp.array([-1, 0, -2, 2]))),
]),
("TupleOfListsOfPythonInts",
[IndexSpec(shape=(3, 4, 5), indexer=([0, 1])),
IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]], [[2, 3, 0]])),
]),
("TupleOfPythonIntsAndIntArrays",
[IndexSpec(shape=(3, 4, 5), indexer=(0, onp.array([0, 1]))),
IndexSpec(shape=(3, 4, 5), indexer=(0, 1,
onp.array([[2, 3, 0]]))),
]),
("TupleOfListsOfPythonIntsAndIntArrays",
[IndexSpec(shape=(3, 4, 5), indexer=([0, 1], onp.array([0]))),
IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]],
onp.array([[2, 3, 0]]))),
]),
]
MIXED_ADVANCED_INDEXING_TESTS_NO_REPEATS = [
("SlicesAndOneIntArrayIndex",
[IndexSpec(shape=(2, 3), indexer=(onp.array([0, 1]), slice(1, 2))),
IndexSpec(shape=(2, 3), indexer=(slice(0, 2),
onp.array([0, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
onp.array([0, 2]),
slice(None))),
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
onp.array([[0, 2], [1, 3]]),
slice(None))),
]),
("SlicesAndTwoIntArrayIndices",
[IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
onp.array([0, 2]),
onp.array([-1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
Ellipsis,
onp.array([-1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
onp.array([-1, 2]),
Ellipsis)),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
onp.array([-1, 2]),
slice(1, 3))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
slice(1, 3),
onp.array([-1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2, -2]),
slice(None, None, 2),
onp.array([-1, 2, 1]))),
]),
("NonesAndIntArrayIndices",
[IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
None,
onp.array([-1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
None,
None,
onp.array([-1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
onp.array([0, 2]),
None,
None,
onp.array([-1, 2]))),
]),
("IntArrayWithInt32Type",
[IndexSpec(shape=(3, 4), indexer=(Ellipsis, onp.array(1, dtype=onp.int32)))
]),
]
MIXED_ADVANCED_INDEXING_TESTS = MIXED_ADVANCED_INDEXING_TESTS_NO_REPEATS + [
("SlicesAndOneIntArrayIndex",
[
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
onp.array([[0, 2], [1, 1]]),
slice(None))),
]),
("SlicesAndTwoIntArrayIndices",
[IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2, -2]),
slice(None, None, 2),
onp.array([-1, 2, -1]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([[0, 2], [2, 0]]),
Ellipsis,
onp.array([[1, 0], [1, 0]]))),
]),]
def dynamic_slice_reference(operand, start_indices, slice_sizes):
out = onp.zeros(slice_sizes, dtype=operand.dtype)
idx = tuple(slice(start, start+size)
for start, size in zip(start_indices, slice_sizes))
section = operand[idx]
out[tuple(slice(None, stop) for stop in section.shape)] = section
return out
def dynamic_update_slice_reference(operand, update, start_indices):
slices = tuple(map(
slice, start_indices, onp.add(start_indices, update.shape)))
updated_operand = onp.copy(operand)
updated_operand[slices] = update
return updated_operand
class IndexingTest(jtu.TestCase):
"""Tests for Numpy indexing translation rules."""
@parameterized.named_parameters(jtu.cases_from_list({
"testcase_name": "{}_inshape={}_indexer={}".format(
name, jtu.format_shape_dtype_string( shape, dtype), indexer),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory, "indexer": indexer
} for name, index_specs in STATIC_INDEXING_TESTS
for shape, indexer in index_specs
for dtype in all_dtypes
for rng_factory in [jtu.rand_default]))
def testStaticIndexing(self, shape, dtype, rng_factory, indexer):
# TODO(rohanj): Revisit passing in self.rng() to this to customize further.
# This would need updating lax_numpy_test as well.
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype)]
onp_fun = lambda x: x[indexer]
jnp_fun = lambda x: tnp.asarray(x)[indexer]
self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True,
check_incomplete_shape=True)
def _ReplaceSlicesWithTuples(self, idx):
"""Helper method to replace slices with tuples for dynamic indexing args."""
if isinstance(idx, slice):
triple = idx.start, idx.stop, idx.step
isnone = [i for i, elt in enumerate(triple) if elt is None]
zeros = itertools.repeat(0)
nones = itertools.repeat(None)
out = subvals(triple, zip(isnone, zeros))
return out, lambda out: slice(*subvals(out, zip(isnone, nones)))
elif isinstance(idx, (tuple, list)) and idx:
t = type(idx)
elts, packs = zip(*map(self._ReplaceSlicesWithTuples, idx))
return elts, lambda elts: t((pack(i) for pack, i in zip(packs, elts)))
else:
return idx, lambda x: x
@parameterized.named_parameters(
{"testcase_name": "{}_inshape={}_indexer={}"
.format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory, "indexer": indexer}
for name, index_specs in [
("OneSliceIndex",
[IndexSpec(shape=(5,), indexer=slice(1, 3)),
IndexSpec(shape=(5, 4), indexer=slice(1, 3))]),
("TwoSliceIndices",
[IndexSpec(shape=(5, 4), indexer=(slice(1, 3), slice(0, 2))),
IndexSpec(shape=(5, 4, 3), indexer=(slice(1, 3), slice(0, 2)))]),
("NonUnitStrides", [
IndexSpec(shape=(3,), indexer=slice(None, None, -1)),
IndexSpec(shape=(3, 3), indexer=slice(0, 3, -2)),
IndexSpec(shape=(3, 4, 5), indexer=slice(0, 4, 2))
]),
("OnlyStartOrStopDynamic", [
IndexSpec(shape=(5, 4), indexer=(slice(None, 3), slice(0, 2))),
IndexSpec(shape=(5, 4, 3), indexer=(slice(1, 3), slice(0, None)))
]),
]
for shape, indexer in index_specs
for dtype in all_dtypes
for rng_factory in [jtu.rand_default])
def testDynamicIndexingWithSlices(self, shape, dtype, rng_factory, indexer):
rng = rng_factory()
unpacked_indexer, pack_indexer = self._ReplaceSlicesWithTuples(indexer)
def onp_fun(x, unpacked_indexer):
indexer = pack_indexer(unpacked_indexer)
return x[indexer]
jnp_fun = lambda x, idx: onp_fun(tnp.asarray(x), idx)
args_maker = lambda: [rng(shape, dtype), unpacked_indexer]
self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
# TODO(wangpeng): check_xla_forced_compile is turned off because some
# compile-time-constant requirements are violated. Investigate and turn it
# on.
self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True,
check_eval_on_shapes=False,
check_incomplete_shape=True,
check_xla_forced_compile=False)
@parameterized.named_parameters(
{"testcase_name": "{}_inshape={}_indexer={}"
.format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory, "indexer": indexer}
for name, index_specs in [
("OneIntIndex",
[IndexSpec(shape=(3,), indexer=1),
IndexSpec(shape=(3, 3), indexer=0),
IndexSpec(shape=(3, 4, 5), indexer=2),
IndexSpec(shape=(3,), indexer=-1),
IndexSpec(shape=(3,), indexer=-2)]),
("TwoIntIndices",
[IndexSpec(shape=(3, 3), indexer=(2, 1)),
IndexSpec(shape=(3, 4, 5), indexer=(1, 2)),
IndexSpec(shape=(3, 4, 5), indexer=(-1, 2))]),
("ThreeIntIndices",
[IndexSpec((3, 4, 5), indexer=(1, 2, 3))]),
]
for shape, indexer in index_specs
for dtype in all_dtypes
for rng_factory in [jtu.rand_default])
def testDynamicIndexingWithIntegers(self, shape, dtype, rng_factory, indexer):
# TODO(rohanj): Revisit passing in self.rng() to this to customize further.
# This would need updating lax_numpy_test as well.
rng = rng_factory()
unpacked_indexer, pack_indexer = self._ReplaceSlicesWithTuples(indexer)
def onp_fun(x, unpacked_indexer):
indexer = pack_indexer(unpacked_indexer)
return x[indexer]
jnp_fun = lambda x, idx: onp_fun(tnp.asarray(x), idx)
args_maker = lambda: [rng(shape, dtype), unpacked_indexer]
self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True,
check_incomplete_shape=True)
@parameterized.named_parameters(
{"testcase_name": "_{}_inshape={}_indexer={}" # pylint: disable=g-complex-comprehension
.format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
"name": name, "shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer}
for name, index_specs in ADVANCED_INDEXING_TESTS
for shape, indexer in index_specs
for dtype in all_dtypes
for rng_factory in [jtu.rand_default])
def testAdvancedIntegerIndexing(self, name, shape, dtype, rng_factory,
indexer):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), indexer]
onp_fun = lambda x, idx: x[idx]
jnp_fun = lambda x, idx: onp_fun(tnp.asarray(x), idx)
self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
# TODO(wangpeng): check_xla_forced_compile is turned off for
# ListOfPythonIntsAndIntArrays because it throws "The number of output
# elements has to equal to number of input elements that are sliced when
# input indices are not constant". Investigate and turn it on.
check_xla = (name != "ListOfPythonIntsAndIntArrays")
self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True,
check_incomplete_shape=True,
check_xla_forced_compile=check_xla)
@parameterized.named_parameters(
{"testcase_name": "_{}_inshape={}_indexer={}" # pylint: disable=g-complex-comprehension
.format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
"name": name, "shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer}
for name, index_specs in MIXED_ADVANCED_INDEXING_TESTS
for shape, indexer in index_specs
for dtype in all_dtypes
for rng_factory in [jtu.rand_default])
def testMixedAdvancedIntegerIndexing(self, name, shape, dtype, rng_factory,
indexer):
rng = rng_factory()
indexer_with_dummies = [e if isinstance(e, onp.ndarray) else ()
for e in indexer]
substitutes = [(i, e) for i, e in enumerate(indexer)
if not isinstance(e, onp.ndarray)]
args_maker = lambda: [rng(shape, dtype), indexer_with_dummies]
def np_fun(x, indexer_with_dummies):
idx = type(indexer)(subvals(indexer_with_dummies, substitutes))
return x[idx]
jnp_fun = lambda x, idx: np_fun(tnp.asarray(x), idx)
self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=True)
# TODO(wangpeng): check_xla_forced_compile is turned off for
# IntArrayWithInt32Type because it throws "The number of output elements has
# to equal to number of input elements that are sliced when input indices
# are not constant". Investigate and turn it on.
check_xla = (name != "IntArrayWithInt32Type")
self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True,
check_incomplete_shape=True,
check_xla_forced_compile=check_xla)
def testAdvancedIndexingManually(self):
x = onp.random.RandomState(0).randn(3, 4, 5)
index_array = onp.array([0, 2, -1, 0])
op = lambda x, index_array: x[..., index_array, :]
cop = nje.jit(op)
a1 = op(x, index_array)
a2 = cop(x, index_array)
self.assertAllClose(a1, a2, check_dtypes=True)
op = lambda x, index_array: x[..., index_array, :, index_array, None]
cop = nje.jit(op)
a1 = op(x, index_array)
a2 = cop(x, index_array)
self.assertAllClose(a1, a2, check_dtypes=True)
op = lambda x, index_array: x[index_array, ..., index_array[:, None], None]
cop = nje.jit(op)
a1 = op(x, index_array)
a2 = cop(x, index_array)
self.assertAllClose(a1, a2, check_dtypes=True)
# Note that we don't currently allow __iter__ in graph mode. So this test only
# iterates over eager tensor.
def testUnpacking(self):
def foo(x):
a, b, c = x
return a + b + c
a1 = foo(onp.arange(3))
a2 = foo(tnp.arange(3))
self.assertAllClose(a1, a2, check_dtypes=True)
def testBooleanIndexingArray1D(self):
idx = onp.array([True, True, False])
x = tnp.asarray(onp.arange(3))
ans = x[idx]
expected = onp.arange(3)[idx]
self.assertAllClose(ans, expected, check_dtypes=False)
def testBooleanIndexingList1D(self):
idx = [True, True, False]
x = tnp.asarray(onp.arange(3))
ans = x[idx]
expected = onp.arange(3)[idx]
self.assertAllClose(ans, expected, check_dtypes=False)
def testBooleanIndexingArray2DBroadcast(self):
idx = onp.array([True, True, False, True])
x = onp.arange(8).reshape(4, 2)
ans = tnp.asarray(x)[idx]
expected = x[idx]
self.assertAllClose(ans, expected, check_dtypes=False)
def testBooleanIndexingList2DBroadcast(self):
idx = [True, True, False, True]
x = onp.arange(8).reshape(4, 2)
ans = tnp.asarray(x)[idx]
expected = x[idx]
self.assertAllClose(ans, expected, check_dtypes=False)
def testBooleanIndexingArray2D(self):
idx = onp.array([[True, False],
[False, True],
[False, False],
[True, True]])
x = onp.arange(8).reshape(4, 2)
ans = tnp.asarray(x)[idx]
expected = x[idx]
self.assertAllClose(ans, expected, check_dtypes=False)
def testBooleanIndexingDynamicShape(self):
x = onp.zeros(3)
i = onp.array([True, True, False])
ans = x[i]
expected = tnp.asarray(x)[i]
self.assertAllClose(ans, expected, check_dtypes=True)
def testIssue187(self):
x = tnp.ones((5, 5))
x[[0, 2, 4], [0, 2, 4]] # doesn't crash
x = onp.arange(25).reshape((5, 5))
ans = nje.jit(lambda x: x[[0, 2, 4], [0, 2, 4]])(x)
expected = x[[0, 2, 4], [0, 2, 4]]
self.assertAllClose(ans, expected, check_dtypes=False)
# TODO(agarwal): Fix this use case.
@jtu.disable
def testIndexingEmptyDimension(self):
# Issue 2671: XLA error when indexing into dimension of size 0
x = tnp.ones((2, 0))
# The following work, even on axis 1 of size 0
_ = x[0, :] + x[0, None] + x[0, 1:] + x[0, 1:3:2]
with self.assertRaisesRegex(IndexError,
"index .* is out of bounds for axis .* with size 0"):
_ = onp.ones((2, 0))[0, 0] # The numpy error
with self.assertRaisesRegex(IndexError,
"index is out of bounds for axis .* with size 0"):
_ = x[0, 0] # JAX indexing
with self.assertRaisesRegex(IndexError,
"index is out of bounds for axis .* with size 0"):
nje.jit(lambda i: x[0, i])(0) # JAX indexing under jit
def testBooleanIndexingWithEmptyResult(self):
# based on a TensorFlow Probability test that started failing after #1623
x = tnp.array([-1])
mask = tnp.array([False])
ans = x[mask] # doesn't crash
expected = onp.array([-1])[onp.array([False])]
self.assertAllClose(ans, expected, check_dtypes=False)
def testFloatIndexingError(self):
error_regex = "only integers, slices.*are valid indices"
# Verify onp behavior
with self.assertRaisesRegex(IndexError, error_regex):
_ = onp.zeros((2, 2))[(0, 0.)]
# Test tnp
with self.assertRaisesRegex(IndexError, error_regex):
tnp.zeros(2)[0.] # pylint: disable=expression-not-assigned
with self.assertRaisesRegex(IndexError, error_regex):
tnp.zeros((2, 2))[(0, 0.)] # pylint: disable=expression-not-assigned
# Test with jit
with self.assertRaisesRegex(IndexError, error_regex):
nje.jit(lambda idx: tnp.zeros((2, 2))[idx])((0, 0.0))
def testIndexOutOfBounds(self): # https://github.com/google/jax/issues/2245
array = tnp.ones(5)
self.assertAllClose(array, array[:10], check_dtypes=True)
@parameterized.named_parameters(jtu.cases_from_list(
{"testcase_name": "_shape={}_start_indices={}_size_indices={}".format( # pylint: disable=g-complex-comprehension
jtu.format_shape_dtype_string(shape, dtype),
start_indices, size_indices),
"shape": shape, "dtype": dtype, "start_indices": start_indices,
"size_indices": size_indices, "rng_factory": rng_factory}
for shape, start_indices, size_indices in [
[(3,), onp.array((1,)), (1,)],
[(5, 3), (1, 1), (3, 1)],
[(5, 3), (1, -2), (3, 1)],
[(5, 3), onp.array((1, 1)), (3, 1)],
[(7, 5, 3), onp.array((4, 1, 0)), (2, 0, 1)],
[(), (), ()],
]
for dtype in default_dtypes
for rng_factory in [jtu.rand_default]))
def testDynamicSlice(self, shape, dtype, start_indices, size_indices,
rng_factory):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), onp.array(start_indices)]
op = lambda x, starts: nje.dynamic_slice(x, starts, size_indices)
self._CompileAndCheck(op, args_maker)
@parameterized.named_parameters(jtu.cases_from_list(
{"testcase_name": "_shape={}_start_indices={}_size_indices={}".format( # pylint: disable=g-complex-comprehension
jtu.format_shape_dtype_string(shape, dtype),
start_indices, size_indices),
"shape": shape, "dtype": dtype, "start_indices": start_indices,
"size_indices": size_indices, "rng_factory": rng_factory}
for shape, start_indices, size_indices in [
[(3,), (1,), (1,)],
[(5, 3), (1, 1), (3, 1)],
[(5, 3), (1, -2), (3, 1)],
[(7, 5, 3), (4, 1, 0), (2, 0, 1)],
[(), (), ()],
]
for dtype in default_dtypes
for rng_factory in [jtu.rand_default]))
def testDynamicSliceAgainstNumpy(self, shape, dtype, start_indices,
size_indices, rng_factory):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), onp.array(start_indices)]
op = lambda x, s: nje.dynamic_slice(x, s, size_indices)
numpy_op = lambda x, s: dynamic_slice_reference(x, s, size_indices)
self._CheckAgainstNumpy(numpy_op, op, args_maker)
def testDynamicSliceInDim(self):
rng = jtu.rand_default()
x = rng((6, 7), onp.int32)
self.assertAllClose(
nje.dynamic_slice_in_dim(x, 2, 3), x[2:5], check_dtypes=True
)
def _broadcastable_shapes(shape):
"""Returns all shapes that broadcast to `shape`."""
def f(rshape):
yield []
if rshape:
for s in f(rshape[1:]):
yield rshape[0:1] + s
if rshape[0] != 1:
for s in f(rshape[1:]):
yield [1] + s
for x in f(list(reversed(shape))):
yield list(reversed(x))
def _update_shape(shape, indexer):
return onp.zeros(shape)[indexer].shape
class UpdateOps(enum.Enum):
UPDATE = 0
ADD = 1
# MUL = 2
MIN = 3
MAX = 4
def np_fn(op, indexer, x, y): # pylint: disable=no-self-argument
x = x.copy()
x[indexer] = {
UpdateOps.UPDATE: lambda: y,
UpdateOps.ADD: lambda: x[indexer] + y,
# UpdateOps.MUL: lambda: x[indexer] * y,
UpdateOps.MIN: lambda: onp.minimum(x[indexer], y),
UpdateOps.MAX: lambda: onp.maximum(x[indexer], y),
}[op]()
return x
def tfnp_fn(op, indexer, x, y): # pylint: disable=no-self-argument
return {
UpdateOps.UPDATE: nje.index_update,
UpdateOps.ADD: nje.index_add,
# UpdateOps.MUL: nje.index_mul,
UpdateOps.MIN: nje.index_min,
UpdateOps.MAX: nje.index_max,
}[op](x, indexer, y)
# a test to workaround b/123559667
def has_non_trivial_stride(indexer):
def has(idx):
return isinstance(idx, slice) and idx.step not in (1, -1, None)
return any(has(idx) for idx in nest.flatten(indexer))
class IndexedUpdateTest(jtu.TestCase):
@parameterized.named_parameters(jtu.cases_from_list({ # pylint: disable=g-complex-comprehension
"testcase_name": "_{}_{}_{}_{}".format(
jtu.format_shape_dtype_string(shape, dtype), indexer,
jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer, "update_shape": update_shape,
"update_dtype": update_dtype, "op": op
} for name, index_specs in STATIC_INDEXING_TESTS
for shape, indexer in index_specs
for op in UpdateOps
for dtype in (all_dtypes if op == UpdateOps.UPDATE else default_dtypes)
for update_shape in _broadcastable_shapes(_update_shape(shape, indexer))
for update_dtype in all_dtypes
for rng_factory in [jtu.rand_default]))
def testStaticIndexing(self, shape, dtype, update_shape, update_dtype,
rng_factory, indexer, op):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
tfnp_fn = lambda x, y: UpdateOps.tfnp_fn(op, indexer, x, y)
self._CheckAgainstNumpy(np_fn, tfnp_fn, args_maker)
# TODO(wangpeng): When indexer is slice(_, 8, -1), XLA throws error "Missing
# xla_context 0-th output from". Investigate.
check_xla = (not has_non_trivial_stride(indexer) and # b/123559667
not (isinstance(indexer, slice) and indexer.stop == 8 and
indexer.step == -1))
self._CompileAndCheck(tfnp_fn, args_maker, check_incomplete_shape=True,
check_experimental_compile=check_xla,
check_xla_forced_compile=check_xla)
@parameterized.named_parameters(jtu.cases_from_list({ # pylint: disable=g-complex-comprehension
"testcase_name": "_{}_{}_{}_{}".format(
jtu.format_shape_dtype_string(shape, dtype), indexer,
jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer, "update_shape": update_shape,
"update_dtype": update_dtype, "op": op
} for name, index_specs in ADVANCED_INDEXING_TESTS_NO_REPEATS
for shape, indexer in index_specs
for op in UpdateOps
for dtype in (all_dtypes if op == UpdateOps.UPDATE else default_dtypes)
for update_shape in _broadcastable_shapes(_update_shape(shape, indexer))
for update_dtype in all_dtypes
for rng_factory in [jtu.rand_default]))
def testAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
rng_factory, indexer, op):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
tfnp_fn = lambda x, y: UpdateOps.tfnp_fn(op, indexer, x, y)
self._CheckAgainstNumpy(np_fn, tfnp_fn, args_maker)
self._CompileAndCheck(tfnp_fn, args_maker, check_incomplete_shape=True)
@parameterized.named_parameters(jtu.cases_from_list({ # pylint: disable=g-complex-comprehension
"testcase_name": "_{}_{}_{}_{}".format(
jtu.format_shape_dtype_string(shape, dtype), indexer,
jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer, "update_shape": update_shape,
"update_dtype": update_dtype, "op": op
} for name, index_specs in MIXED_ADVANCED_INDEXING_TESTS_NO_REPEATS
for shape, indexer in index_specs
for op in UpdateOps
for dtype in (all_dtypes if op == UpdateOps.UPDATE else default_dtypes)
for update_shape in _broadcastable_shapes(_update_shape(shape, indexer))
for update_dtype in all_dtypes
for rng_factory in [jtu.rand_default]))
def testMixedAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
rng_factory, indexer, op):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
tfnp_fn = lambda x, y: UpdateOps.tfnp_fn(op, indexer, x, y)
self._CheckAgainstNumpy(np_fn, tfnp_fn, args_maker)
check_xla = not has_non_trivial_stride(indexer) # b/123559667
self._CompileAndCheck(tfnp_fn, args_maker, check_incomplete_shape=True,
check_experimental_compile=check_xla,
check_xla_forced_compile=check_xla)
@parameterized.named_parameters(jtu.cases_from_list({ # pylint: disable=g-complex-comprehension
"testcase_name": "_{}_{}_{}_{}".format(
jtu.format_shape_dtype_string(shape, dtype), indexer,
jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer, "update_shape": update_shape,
"update_dtype": update_dtype, "op": op
} for name, index_specs in STATIC_INDEXING_TESTS
for shape, indexer in index_specs
for op in [UpdateOps.ADD, UpdateOps.UPDATE]
for dtype in float_dtypes
for update_shape in _broadcastable_shapes(_update_shape(shape, indexer))
for update_dtype in float_dtypes
for rng_factory in [jtu.rand_default]))
def testStaticIndexingGrads(self, shape, dtype, update_shape, update_dtype,
rng_factory, indexer, op):
rng = rng_factory()
tfnp_fn = lambda x, y: UpdateOps.tfnp_fn(op, indexer, x, y)
x = rng(shape, dtype)
y = rng(update_shape, update_dtype)
self.check_grads(tfnp_fn, (x, y), rtol=1e-3, atol=1e-3, delta=1.)
@parameterized.named_parameters(jtu.cases_from_list(
{"testcase_name": "_shape={}_start_indices={}_update_shape={}".format( # pylint: disable=g-complex-comprehension
jtu.format_shape_dtype_string(shape, dtype),
start_indices, update_shape),
"shape": shape, "dtype": dtype, "start_indices": start_indices,
"update_shape": update_shape, "rng_factory": rng_factory}
for shape, start_indices, update_shape in [
[(3,), (1,), (1,)],
[(5, 3), (1, 1), (3, 1)],
[(5, 3), (1, -2), (3, 1)],
[(7, 5, 3), (4, 1, 0), (2, 0, 1)],
[(), (), ()],
]
for dtype in default_dtypes
for rng_factory in [jtu.rand_default]))
def testDynamicUpdateSlice(self, shape, dtype, start_indices, update_shape,
rng_factory):
rng = rng_factory()
def args_maker():
return [rng(shape, dtype), rng(update_shape, dtype),
onp.array(start_indices)]
# update's shape must be fully known.
# TODO(wangpeng): Support turning off check_incomplete_shape for individual
# arguments.
self._CompileAndCheck(
nje.dynamic_update_slice, args_maker, check_incomplete_shape=False
)
@parameterized.named_parameters(jtu.cases_from_list(
{"testcase_name": "_shape={}_start_indices={}_update_shape={}".format( # pylint: disable=g-complex-comprehension
jtu.format_shape_dtype_string(shape, dtype),
start_indices, update_shape),
"shape": shape, "dtype": dtype, "start_indices": start_indices,
"update_shape": update_shape, "rng_factory": rng_factory}
for shape, start_indices, update_shape in [
[(3,), (1,), (1,)],
[(5, 3), (1, 1), (3, 1)],
[(5, 3), (1, -2), (3, 1)],
[(7, 5, 3), (4, 1, 0), (2, 0, 1)],
[(), (), ()],
]
for dtype in default_dtypes
for rng_factory in [jtu.rand_default]))
def testDynamicUpdateSliceAgainstNumpy(self, shape, dtype, start_indices,
update_shape, rng_factory):
rng = rng_factory()
def args_maker():
return [rng(shape, dtype), rng(update_shape, dtype),
onp.array(start_indices)]
self._CheckAgainstNumpy(
dynamic_update_slice_reference, nje.dynamic_update_slice, args_maker
)
def testDynamicUpdateSliceInDim(self):
rng = jtu.rand_default()
x = rng((6, 7), onp.int32)
y = rng((3, 7), onp.int32)
z = x.copy()
z[2:5] = y
self.assertAllClose(
nje.dynamic_update_slice_in_dim(x, y, 2, 0), z, check_dtypes=True
)
if __name__ == "__main__":
tf_config.set_soft_device_placement(False)
absltest.main()
tensorflow/python/ops/numpy_ops/tests/np_test.py
+3166
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/numpy_ops/tests/np_wrapper.py
+43
View File
@@ -0,0 +1,43 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""TF NumPy API wrapper for the tests."""
# pylint: disable=wildcard-import
# pylint: disable=unused-import
# pylint: disable=g-importing-member
import numpy as onp
from tensorflow.python.compat import v2_compat
from tensorflow.python.framework.dtypes import bfloat16
from tensorflow.python.ops.numpy_ops import np_random as random
from tensorflow.python.ops.numpy_ops.np_array_ops import *
from tensorflow.python.ops.numpy_ops.np_arrays import ndarray
from tensorflow.python.ops.numpy_ops.np_config import enable_numpy_behavior
from tensorflow.python.ops.numpy_ops.np_dtypes import *
from tensorflow.python.ops.numpy_ops.np_dtypes import canonicalize_dtype
from tensorflow.python.ops.numpy_ops.np_dtypes import default_float_type
from tensorflow.python.ops.numpy_ops.np_dtypes import is_allow_float64
from tensorflow.python.ops.numpy_ops.np_dtypes import set_allow_float64
from tensorflow.python.ops.numpy_ops.np_math_ops import *
from tensorflow.python.ops.numpy_ops.np_utils import finfo
from tensorflow.python.ops.numpy_ops.np_utils import promote_types
from tensorflow.python.ops.numpy_ops.np_utils import result_type
random.DEFAULT_RANDN_DTYPE = onp.float32
# pylint: enable=unused-import
v2_compat.enable_v2_behavior()
# TODO(b/171429739): This should be moved to every individual file/test.
enable_numpy_behavior()
tensorflow/python/ops/numpy_ops/tests/test_util.py
+899
View File
@@ -0,0 +1,899 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""NumPy test utilities."""
from contextlib import contextmanager
import functools
from functools import partial
import re
import itertools as it
import os
from typing import Dict, Sequence, Union
import unittest
import warnings
import zlib
from absl.testing import absltest
from absl.testing import parameterized
import numpy as onp
import numpy.random as npr
from tensorflow.python.util import nest
from tensorflow.python.util import numpy_compat
from tensorflow.python.framework import tensor
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import gradient_checker_v2
from tensorflow.python.ops.numpy_ops.tests.config import flags
import tensorflow.python.ops.numpy_ops.tests.extensions as nje
from tensorflow.python.ops.numpy_ops import np_utils
from tensorflow.python.ops.numpy_ops import np_array_ops
tree_map = nest.map_structure
tree_multimap = nest.map_structure
FLAGS = flags.FLAGS
# https://danielms.site/zet/2023/pythons-distutil-strtobool-replacement/
def strtobool(value: str) -> bool:
value = value.lower()
if value in ('y', 'yes', 'on', '1', 'true', 't'):
return True
return False
# TODO(wangpeng): Remove this flag after broken tests are fixed
flags.DEFINE_bool('enable_x64',
strtobool('False'),
'Enable 64-bit types to be used.')
flags.DEFINE_enum(
'test_dut', '',
enum_values=['', 'cpu', 'gpu', 'tpu'],
help=
'Describes the device under test in case special consideration is required.'
)
flags.DEFINE_integer(
'num_generated_cases',
10,
help='Number of generated cases to test')
EPS = 1e-4
# Default dtypes corresponding to Python scalars.
python_scalar_dtypes = {
bool: onp.dtype(onp.bool_),
int: onp.dtype(onp.int_),
float: onp.dtype(onp.float64),
complex: onp.dtype(onp.complex128),
}
def _dtype(x):
if isinstance(x, tensor.Tensor):
return x.dtype.as_numpy_dtype
return (getattr(x, 'dtype', None) or
onp.dtype(python_scalar_dtypes.get(type(x), None)) or
numpy_compat.np_asarray(x).dtype)
def is_sequence(x):
try:
iter(x)
except TypeError:
return False
else:
return True
_default_tolerance = {
onp.dtype(onp.bool_): 0,
onp.dtype(onp.int8): 0,
onp.dtype(onp.int16): 0,
onp.dtype(onp.int32): 0,
onp.dtype(onp.int64): 0,
onp.dtype(onp.uint8): 0,
onp.dtype(onp.uint16): 0,
onp.dtype(onp.uint32): 0,
onp.dtype(onp.uint64): 0,
# TODO(b/154768983): onp.dtype(dtypes.bfloat16): 1e-2,
onp.dtype(onp.float16): 1e-3,
onp.dtype(onp.float32): 1e-6,
onp.dtype(onp.float64): 1e-15,
onp.dtype(onp.complex64): 1e-6,
onp.dtype(onp.complex128): 1e-15,
}
def default_tolerance():
return _default_tolerance
default_gradient_tolerance = {
# TODO(b/154768983): onp.dtype(dtypes.bfloat16): 1e-1,
onp.dtype(onp.float16): 1e-2,
onp.dtype(onp.float32): 2e-3,
onp.dtype(onp.float64): 1e-5,
onp.dtype(onp.complex64): 1e-3,
onp.dtype(onp.complex128): 1e-5,
}
def _assert_numpy_allclose(a, b, atol=None, rtol=None):
# TODO(b/154768983):
# a = a.astype(onp.float32) if a.dtype == dtypes.bfloat16 else a
# b = b.astype(onp.float32) if b.dtype == dtypes.bfloat16 else b
kw = {}
if atol: kw["atol"] = atol
if rtol: kw["rtol"] = rtol
onp.testing.assert_allclose(a, b, **kw)
def tolerance(dtype, tol=None):
tol = {} if tol is None else tol
if not isinstance(tol, dict):
return tol
tol = {onp.dtype(key): value for key, value in tol.items()}
dtype = onp.dtype(dtype)
return tol.get(dtype, default_tolerance()[dtype])
def _normalize_tolerance(tol):
tol = tol or 0
if isinstance(tol, dict):
return {onp.dtype(k): v for k, v in tol.items()}
else:
return {k: tol for k in _default_tolerance}
def join_tolerance(tol1, tol2):
tol1 = _normalize_tolerance(tol1)
tol2 = _normalize_tolerance(tol2)
out = tol1
for k, v in tol2.items():
out[k] = max(v, tol1.get(k, 0))
return out
def _assert_numpy_close(a, b, atol=None, rtol=None):
assert a.shape == b.shape
atol = max(tolerance(a.dtype, atol), tolerance(b.dtype, atol))
rtol = max(tolerance(a.dtype, rtol), tolerance(b.dtype, rtol))
_assert_numpy_allclose(a, b, atol=atol * a.size, rtol=rtol * b.size)
def check_eq(xs, ys):
tree_all(tree_multimap(_assert_numpy_allclose, xs, ys))
def check_close(xs, ys, atol=None, rtol=None):
assert_close = partial(_assert_numpy_close, atol=atol, rtol=rtol)
tree_all(tree_multimap(assert_close, xs, ys))
def inner_prod(xs, ys):
def contract(x, y):
return onp.real(onp.dot(onp.conj(x).reshape(-1), y.reshape(-1)))
return tree_reduce(onp.add, tree_multimap(contract, xs, ys))
add = partial(tree_multimap, lambda x, y: onp.add(x, y, dtype=_dtype(x)))
sub = partial(tree_multimap, lambda x, y: onp.subtract(x, y, dtype=_dtype(x)))
conj = partial(tree_map, lambda x: onp.conj(x, dtype=_dtype(x)))
def scalar_mul(xs, a):
return tree_map(lambda x: onp.multiply(x, a, dtype=_dtype(x)), xs)
def rand_like(rng, x):
shape = onp.shape(x)
dtype = _dtype(x)
randn = lambda: numpy_compat.np_asarray(rng.randn(*shape), dtype=dtype)
if onp.issubdtype(dtype, onp.complexfloating):
return randn() + dtype.type(1.0j) * randn()
else:
return randn()
def numerical_jvp(f, primals, tangents, eps=EPS):
delta = scalar_mul(tangents, eps)
f_pos = f(*add(primals, delta))
f_neg = f(*sub(primals, delta))
return scalar_mul(sub(f_pos, f_neg), 0.5 / eps)
def _merge_tolerance(tol, default):
if tol is None:
return default
if not isinstance(tol, dict):
return tol
out = default.copy()
for k, v in tol.items():
out[onp.dtype(k)] = v
return out
def check_jvp(f, f_jvp, args, atol=None, rtol=None, eps=EPS):
atol = _merge_tolerance(atol, default_gradient_tolerance)
rtol = _merge_tolerance(rtol, default_gradient_tolerance)
rng = onp.random.RandomState(0)
tangent = tree_map(partial(rand_like, rng), args)
v_out, t_out = f_jvp(args, tangent)
v_out_expected = f(*args)
t_out_expected = numerical_jvp(f, args, tangent, eps=eps)
# In principle we should expect exact equality of v_out and v_out_expected,
# but due to nondeterminism especially on GPU (e.g., due to convolution
# autotuning) we only require "close".
check_close(v_out, v_out_expected, atol=atol, rtol=rtol)
check_close(t_out, t_out_expected, atol=atol, rtol=rtol)
def check_vjp(f, f_vjp, args, atol=None, rtol=None, eps=EPS):
atol = _merge_tolerance(atol, default_gradient_tolerance)
rtol = _merge_tolerance(rtol, default_gradient_tolerance)
_rand_like = partial(rand_like, onp.random.RandomState(0))
v_out, vjpfun = f_vjp(*args)
v_out_expected = f(*args)
check_close(v_out, v_out_expected, atol=atol, rtol=rtol)
tangent = tree_map(_rand_like, args)
tangent_out = numerical_jvp(f, args, tangent, eps=eps)
cotangent = tree_map(_rand_like, v_out)
cotangent_out = conj(vjpfun(conj(cotangent)))
ip = inner_prod(tangent, cotangent_out)
ip_expected = inner_prod(tangent_out, cotangent)
check_close(ip, ip_expected, atol=atol, rtol=rtol)
def device_under_test():
return FLAGS.test_dut
def if_device_under_test(device_type: Union[str, Sequence[str]],
if_true, if_false):
"""Chooses `if_true` of `if_false` based on device_under_test."""
if device_under_test() in ([device_type] if isinstance(device_type, str)
else device_type):
return if_true
else:
return if_false
def supported_dtypes():
if device_under_test() == "tpu":
return {onp.bool_, onp.int32, onp.uint32, dtypes.bfloat16, onp.float32,
onp.complex64}
else:
return {onp.bool_, onp.int8, onp.int16, onp.int32, onp.int64,
onp.uint8, onp.uint16, onp.uint32, onp.uint64,
dtypes.bfloat16, onp.float16, onp.float32, onp.float64,
onp.complex64, onp.complex128}
def skip_if_unsupported_type(dtype):
if dtype not in supported_dtypes():
raise unittest.SkipTest(
f"Type {dtype} not supported on {device_under_test()}")
def skip_on_devices(*disabled_devices):
"""A decorator for test methods to skip the test on certain devices."""
def skip(test_method):
@functools.wraps(test_method)
def test_method_wrapper(self, *args, **kwargs):
device = device_under_test()
if device in disabled_devices:
test_name = getattr(test_method, '__name__', '[unknown test]')
raise unittest.SkipTest(
f"{test_name} not supported on {device.upper()}.")
return test_method(self, *args, **kwargs)
return test_method_wrapper
return skip
def skip_on_flag(flag_name, skip_value):
"""A decorator for test methods to skip the test when flags are set."""
def skip(test_method): # pylint: disable=missing-docstring
@functools.wraps(test_method)
def test_method_wrapper(self, *args, **kwargs):
flag_value = getattr(FLAGS, flag_name)
if flag_value == skip_value:
test_name = getattr(test_method, '__name__', '[unknown test]')
raise unittest.SkipTest(
f"{test_name} not supported when FLAGS.{flag_name} is {flag_value}")
return test_method(self, *args, **kwargs)
return test_method_wrapper
return skip
def format_test_name_suffix(opname, shapes, dt):
arg_descriptions = (format_shape_dtype_string(shape, dtype)
for shape, dtype in zip(shapes, dt))
return '{}_{}'.format(opname.capitalize(), '_'.join(arg_descriptions))
# We use special symbols, represented as singleton objects, to distinguish
# between NumPy scalars, Python scalars, and 0-D arrays.
class ScalarShape:
def __len__(self): return 0
def __getitem__(self, i):
raise IndexError(f'index {i} out of range.')
class _NumpyScalar(ScalarShape): pass
class _PythonScalar(ScalarShape): pass
NUMPY_SCALAR_SHAPE = _NumpyScalar()
PYTHON_SCALAR_SHAPE = _PythonScalar()
def _dims_of_shape(shape):
"""Converts `shape` to a tuple of dimensions."""
if type(shape) in (list, tuple):
return shape
elif isinstance(shape, ScalarShape):
return ()
else:
raise TypeError(type(shape))
def _cast_to_shape(value, shape, dtype):
"""Casts `value` to the correct Python type for `shape` and `dtype`."""
if shape is NUMPY_SCALAR_SHAPE:
# explicitly cast to NumPy scalar in case `value` is a Python scalar.
return onp.dtype(dtype).type(value)
elif shape is PYTHON_SCALAR_SHAPE:
# explicitly cast to Python scalar via https://stackoverflow.com/a/11389998
return numpy_compat.np_asarray(value).item()
elif type(shape) in (list, tuple):
assert onp.shape(value) == tuple(shape)
return value
else:
raise TypeError(type(shape))
def dtype_str(dtype):
return onp.dtype(dtype).name
def format_shape_dtype_string(shape, dtype):
if shape is NUMPY_SCALAR_SHAPE:
return dtype_str(dtype)
elif shape is PYTHON_SCALAR_SHAPE:
return 'py' + dtype_str(dtype)
elif type(shape) in (list, tuple):
shapestr = ','.join(str(dim) for dim in shape)
return '{}[{}]'.format(dtype_str(dtype), shapestr)
elif type(shape) is int:
return '{}[{},]'.format(dtype_str(dtype), shape)
elif isinstance(shape, onp.ndarray):
return '{}[{}]'.format(dtype_str(dtype), shape)
else:
raise TypeError(type(shape))
def _rand_dtype(rand, shape, dtype, scale=1., post=lambda x: x):
"""Produce random values given shape, dtype, scale, and post-processor.
Args:
rand: a function for producing random values of a given shape, e.g. a
bound version of either onp.RandomState.randn or onp.RandomState.rand.
shape: a shape value as a tuple of positive integers.
dtype: a numpy dtype.
scale: optional, a multiplicative scale for the random values (default 1).
post: optional, a callable for post-processing the random values (default
identity).
Returns:
An ndarray of the given shape and dtype using random values based on a call
to rand but scaled, converted to the appropriate dtype, and post-processed.
"""
r = lambda: numpy_compat.np_asarray(scale * rand(*_dims_of_shape(shape)),
dtype)
if onp.issubdtype(dtype, onp.complexfloating):
vals = r() + 1.0j * r()
else:
vals = r()
return _cast_to_shape(numpy_compat.np_asarray(post(vals), dtype), shape,
dtype)
def rand_default(scale=3):
randn = npr.RandomState(0).randn
return partial(_rand_dtype, randn, scale=scale)
def rand_nonzero():
post = lambda x: onp.where(x == 0, onp.array(1, dtype=x.dtype), x)
randn = npr.RandomState(0).randn
return partial(_rand_dtype, randn, scale=3, post=post)
def rand_positive():
post = lambda x: x + 1
rand = npr.RandomState(0).rand
return partial(_rand_dtype, rand, scale=2, post=post)
def rand_small():
randn = npr.RandomState(0).randn
return partial(_rand_dtype, randn, scale=1e-3)
def rand_not_small(offset=10.):
post = lambda x: x + onp.where(x > 0, offset, -offset)
randn = npr.RandomState(0).randn
return partial(_rand_dtype, randn, scale=3., post=post)
def rand_small_positive():
rand = npr.RandomState(0).rand
return partial(_rand_dtype, rand, scale=2e-5)
def rand_uniform(low=0.0, high=1.0):
assert low < high
rand = npr.RandomState(0).rand
post = lambda x: x * (high - low) + low
return partial(_rand_dtype, rand, post=post)
def rand_some_equal():
randn = npr.RandomState(0).randn
rng = npr.RandomState(0)
def post(x):
x_ravel = x.ravel()
if len(x_ravel) == 0:
return x
flips = rng.rand(*onp.shape(x)) < 0.5
return onp.where(flips, x_ravel[0], x)
return partial(_rand_dtype, randn, scale=100., post=post)
def rand_some_inf():
"""Return a random sampler that produces infinities in floating types."""
rng = npr.RandomState(1)
base_rand = rand_default()
"""
TODO: Complex numbers are not correctly tested
If blocks should be switched in order, and relevant tests should be fixed
"""
def rand(shape, dtype):
"""The random sampler function."""
if not onp.issubdtype(dtype, onp.floating):
# only float types have inf
return base_rand(shape, dtype)
if onp.issubdtype(dtype, onp.complexfloating):
base_dtype = onp.real(onp.array(0, dtype=dtype)).dtype
out = (rand(shape, base_dtype) +
onp.array(1j, dtype) * rand(shape, base_dtype))
return _cast_to_shape(out, shape, dtype)
dims = _dims_of_shape(shape)
posinf_flips = rng.rand(*dims) < 0.1
neginf_flips = rng.rand(*dims) < 0.1
vals = base_rand(shape, dtype)
vals = onp.where(posinf_flips, onp.array(onp.inf, dtype=dtype), vals)
vals = onp.where(neginf_flips, onp.array(-onp.inf, dtype=dtype), vals)
return _cast_to_shape(numpy_compat.np_asarray(vals, dtype=dtype), shape,
dtype)
return rand
def rand_some_nan():
"""Return a random sampler that produces nans in floating types."""
rng = npr.RandomState(1)
base_rand = rand_default()
def rand(shape, dtype):
"""The random sampler function."""
if onp.issubdtype(dtype, onp.complexfloating):
base_dtype = onp.real(onp.array(0, dtype=dtype)).dtype
out = (rand(shape, base_dtype) +
onp.array(1j, dtype) * rand(shape, base_dtype))
return _cast_to_shape(out, shape, dtype)
if not onp.issubdtype(dtype, onp.floating):
# only float types have inf
return base_rand(shape, dtype)
dims = _dims_of_shape(shape)
nan_flips = rng.rand(*dims) < 0.1
vals = base_rand(shape, dtype)
vals = onp.where(nan_flips, onp.array(onp.nan, dtype=dtype), vals)
return _cast_to_shape(numpy_compat.np_asarray(vals, dtype=dtype), shape,
dtype)
return rand
def rand_some_inf_and_nan():
"""Return a random sampler that produces infinities in floating types."""
rng = npr.RandomState(1)
base_rand = rand_default()
"""
TODO: Complex numbers are not correctly tested
If blocks should be switched in order, and relevant tests should be fixed
"""
def rand(shape, dtype):
"""The random sampler function."""
if not onp.issubdtype(dtype, onp.floating):
# only float types have inf
return base_rand(shape, dtype)
if onp.issubdtype(dtype, onp.complexfloating):
base_dtype = onp.real(onp.array(0, dtype=dtype)).dtype
out = (rand(shape, base_dtype) +
onp.array(1j, dtype) * rand(shape, base_dtype))
return _cast_to_shape(out, shape, dtype)
dims = _dims_of_shape(shape)
posinf_flips = rng.rand(*dims) < 0.1
neginf_flips = rng.rand(*dims) < 0.1
nan_flips = rng.rand(*dims) < 0.1
vals = base_rand(shape, dtype)
vals = onp.where(posinf_flips, onp.array(onp.inf, dtype=dtype), vals)
vals = onp.where(neginf_flips, onp.array(-onp.inf, dtype=dtype), vals)
vals = onp.where(nan_flips, onp.array(onp.nan, dtype=dtype), vals)
return _cast_to_shape(numpy_compat.np_asarray(vals, dtype=dtype), shape,
dtype)
return rand
# TODO(mattjj): doesn't handle complex types
def rand_some_zero():
"""Return a random sampler that produces some zeros."""
rng = npr.RandomState(1)
base_rand = rand_default()
def rand(shape, dtype):
"""The random sampler function."""
dims = _dims_of_shape(shape)
zeros = rng.rand(*dims) < 0.5
vals = base_rand(shape, dtype)
vals = onp.where(zeros, onp.array(0, dtype=dtype), vals)
return _cast_to_shape(numpy_compat.np_asarray(vals, dtype=dtype), shape,
dtype)
return rand
def rand_int(low, high=None):
randint = npr.RandomState(0).randint
def fn(shape, dtype):
return randint(low, high=high, size=shape, dtype=dtype)
return fn
def rand_unique_int():
randchoice = npr.RandomState(0).choice
def fn(shape, dtype):
return randchoice(onp.arange(onp.prod(shape), dtype=dtype),
size=shape, replace=False)
return fn
def rand_bool():
rng = npr.RandomState(0)
def generator(shape, dtype):
return _cast_to_shape(rng.rand(*_dims_of_shape(shape)) < 0.5, shape, dtype)
return generator
def check_raises(thunk, err_type, msg):
try:
thunk()
assert False
except err_type as e:
assert str(e).startswith(msg), "\n{}\n\n{}\n".format(e, msg)
def check_raises_regexp(thunk, err_type, pattern):
try:
thunk()
assert False
except err_type as e:
assert re.match(pattern, str(e)), "{}\n\n{}\n".format(e, pattern)
def _iter_eqns(jaxpr):
# TODO(necula): why doesn't this search in params?
for eqn in jaxpr.eqns:
yield eqn
for subjaxpr in core.subjaxprs(jaxpr):
yield from _iter_eqns(subjaxpr)
def assert_dot_precision(expected_precision, fun, *args):
jaxpr = api.make_jaxpr(fun)(*args)
precisions = [eqn.params['precision'] for eqn in _iter_eqns(jaxpr.jaxpr)
if eqn.primitive == lax.dot_general_p]
for precision in precisions:
msg = "Unexpected precision: {} != {}".format(expected_precision, precision)
assert precision == expected_precision, msg
_CACHED_INDICES: Dict[int, Sequence[int]] = {}
def cases_from_list(xs):
xs = list(xs)
n = len(xs)
k = min(n, FLAGS.num_generated_cases)
# Random sampling for every parameterized test is expensive. Do it once and
# cache the result.
indices = _CACHED_INDICES.get(n)
if indices is None:
rng = npr.RandomState(42)
_CACHED_INDICES[n] = indices = rng.permutation(n)
return [xs[i] for i in indices[:k]]
def cases_from_gens(*gens):
sizes = [1, 3, 10]
cases_per_size = int(FLAGS.num_generated_cases / len(sizes)) + 1
for size in sizes:
for i in range(cases_per_size):
yield ('_{}_{}'.format(size, i),) + tuple(gen(size) for gen in gens)
def to_np(a):
return nest.map_structure(np_array_ops.asarray, a)
def to_tf_fn(f):
return lambda *args: f(*to_np(args))
class TestCase(parameterized.TestCase):
"""Base class for tests including numerical checks and boilerplate."""
# copied from jax.test_util
def setUp(self):
super().setUp()
self._rng = npr.RandomState(zlib.adler32(self._testMethodName.encode()))
# copied from jax.test_util
def rng(self):
return self._rng
# TODO(mattjj): this obscures the error messages from failures, figure out how
# to re-enable it
# def tearDown(self) -> None:
# assert core.reset_trace_state()
def assertArraysAllClose(self, x, y, check_dtypes, atol=None, rtol=None):
"""Assert that x and y are close (up to numerical tolerances)."""
self.assertEqual(x.shape, y.shape)
atol = max(tolerance(_dtype(x), atol), tolerance(_dtype(y), atol))
rtol = max(tolerance(_dtype(x), rtol), tolerance(_dtype(y), rtol))
_assert_numpy_allclose(x, y, atol=atol, rtol=rtol)
if check_dtypes:
self.assertDtypesMatch(x, y)
def assertDtypesMatch(self, x, y):
if FLAGS.enable_x64:
self.assertEqual(_dtype(x), _dtype(y))
def assertAllClose(self, x, y, check_dtypes, atol=None, rtol=None):
"""Assert that x and y, either arrays or nested tuples/lists, are close."""
if isinstance(x, dict):
self.assertIsInstance(y, dict)
self.assertEqual(set(x.keys()), set(y.keys()))
for k in x:
self.assertAllClose(x[k], y[k], check_dtypes, atol=atol, rtol=rtol)
elif is_sequence(x) and not hasattr(x, '__array__'):
self.assertTrue(is_sequence(y) and not hasattr(y, '__array__'))
self.assertEqual(len(x), len(y))
for x_elt, y_elt in zip(x, y):
self.assertAllClose(x_elt, y_elt, check_dtypes, atol=atol, rtol=rtol)
elif hasattr(x, '__array__') or onp.isscalar(x):
self.assertTrue(hasattr(y, '__array__') or onp.isscalar(y))
if check_dtypes:
self.assertDtypesMatch(x, y)
x = numpy_compat.np_asarray(x)
y = numpy_compat.np_asarray(y)
self.assertArraysAllClose(x, y, check_dtypes=False, atol=atol, rtol=rtol)
elif x == y:
return
else:
raise TypeError((type(x), type(y)))
def assertMultiLineStrippedEqual(self, expected, what):
"""Asserts two strings are equal, after stripping each line."""
ignore_space_re = re.compile(r'\s*\n\s*')
expected_clean = re.sub(ignore_space_re, '\n', expected.strip())
what_clean = re.sub(ignore_space_re, '\n', what.strip())
self.assertMultiLineEqual(expected_clean, what_clean,
msg="Found\n{}\nExpecting\n{}".format(what, expected))
def _CheckAgainstNumpy(self, numpy_reference_op, lax_op, args_maker,
check_dtypes=True, tol=None):
args = args_maker()
lax_ans = lax_op(*args)
numpy_ans = numpy_reference_op(*args)
self.assertAllClose(numpy_ans, lax_ans, check_dtypes=check_dtypes,
atol=tol, rtol=tol)
def _CompileAndCheck(self,
fun,
args_maker,
check_dtypes=True,
rtol=None,
atol=None,
check_eval_on_shapes=True,
check_incomplete_shape=True,
check_unknown_rank=True,
static_argnums=(),
check_experimental_compile=True,
check_xla_forced_compile=True):
"""Compiles the function and checks the results.
Args:
fun: the function to be checked.
args_maker: a callable that returns a tuple which will be used as the
positional arguments.
check_dtypes: whether to check that the result dtypes from non-compiled
and compiled runs agree.
rtol: relative tolerance for allclose assertions.
atol: absolute tolerance for allclose assertions.
check_eval_on_shapes: whether to run `eval_on_shapes` on the function and
check that the result shapes and dtypes are correct.
check_incomplete_shape: whether to check that the function can handle
incomplete shapes (including those with and without a known rank).
check_unknown_rank: (only has effect when check_incomplete_shape is True)
whether to check that the function can handle unknown ranks.
static_argnums: indices of arguments to be treated as static arguments for
`jit` and `eval_on_shapes`.
check_experimental_compile: whether to check compilation with
experimental_compile=True (in addition to compilation without the flag).
check_xla_forced_compile: whether to check compilation with
forced_compile=True (in addition to compilation without the flag). This
flag is different from experimental_compile because it enforces
whole-function compilation while the latter doesn't. TPU requires
whole-function compilation.
"""
args = args_maker()
for x in args:
if not hasattr(x, 'dtype'):
# If there is a input that doesn't have dtype info, jit and
# eval_on_shapes may pick a different dtype for it than numpy, so we
# skip the dtype check.
check_dtypes = False
python_ans = fun(*args)
python_shapes = nest.map_structure(onp.shape, python_ans)
onp_shapes = nest.map_structure(
lambda x: onp.shape(numpy_compat.np_asarray(x)), python_ans
)
self.assertEqual(python_shapes, onp_shapes)
def check_compile(**kwargs):
# `wrapped_fun` and `python_should_be_executing` are used to check that
# when the jitted function is called the second time, the original Python
# function won't be executed.
def wrapped_fun(*args):
self.assertTrue(python_should_be_executing)
return fun(*args)
cfun = nje.jit(wrapped_fun, static_argnums=static_argnums, **kwargs)
python_should_be_executing = True
monitored_ans = cfun(*args)
python_should_be_executing = False
compiled_ans = cfun(*args)
self.assertAllClose(python_ans, monitored_ans, check_dtypes, atol, rtol)
self.assertAllClose(python_ans, compiled_ans, check_dtypes, atol, rtol)
# Run `cfun` with a different set of arguments to check that changing
# arguments won't cause recompilation.
new_args = args_maker()
skip_retracing_test = False
for old, new in zip(nest.flatten(args), nest.flatten(new_args)):
if nje.most_precise_int_dtype(old) != nje.most_precise_int_dtype(new):
# If the old and new arguments result in different dtypes (because
# they fall into different value ranges), tf-numpy will retrace, so we
# skip the no-retrace test.
skip_retracing_test = True
if not skip_retracing_test:
python_should_be_executing = True
new_python_ans = fun(*new_args)
python_should_be_executing = False
compiled_ans = cfun(*new_args)
self.assertAllClose(new_python_ans, compiled_ans, check_dtypes, atol,
rtol)
check_compile()
if check_experimental_compile:
check_compile(experimental_compile=True)
if check_xla_forced_compile:
check_compile(xla_forced_compile=True)
if check_eval_on_shapes:
# Check that nje.eval_on_shapes can get complete output shapes given
# complete input shapes.
cfun = nje.eval_on_shapes(fun, static_argnums=static_argnums)
compiled_ans = cfun(*args)
flat_python_ans = nest.flatten(python_ans)
flat_compiled_ans = nest.flatten(compiled_ans)
self.assertEqual(len(flat_python_ans), len(flat_compiled_ans))
for a, b in zip(flat_python_ans, flat_compiled_ans):
if hasattr(a, 'shape'):
self.assertEqual(a.shape, b.shape)
if check_dtypes and hasattr(a, 'dtype'):
self.assertEqual(dtypes.as_dtype(a.dtype), b.dtype)
# If some argument doesn't have a `dtype` attr (e.g. a Python scalar), we
# skip incomplete-shape checks, since shape specs need dtype. It's OK to
# skip since the same incomplete-shape checks will run for []-shaped arrays.
if check_incomplete_shape and all(hasattr(x, 'dtype') for x in args):
# Check partial shapes with known ranks.
# Numpy scalars (created by e.g. np.int32(5)) have `dtype` but not
# `shape`.
if all(hasattr(x, 'shape') for x in args):
specs = [
tensor.TensorSpec([None] * len(x.shape), x.dtype) for x in args
]
cfun = nje.jit(
fun, static_argnums=static_argnums, input_signature=specs
)
compiled_ans = cfun(*args)
self.assertAllClose(python_ans, compiled_ans, check_dtypes, atol, rtol)
if check_unknown_rank:
# Check unknown ranks.
specs = [tensor.TensorSpec(None, x.dtype) for x in args]
cfun = nje.jit(
fun, static_argnums=static_argnums, input_signature=specs)
compiled_ans = cfun(*args)
self.assertAllClose(python_ans, compiled_ans, check_dtypes, atol, rtol)
def check_grads(self, f, args, atol=None, rtol=None, delta=None):
"""Check gradients against finite differences.
Args:
f: function to check at ``f(*args)``.
args: a list or tuple of argument values.
atol: absolute tolerance for gradient equality.
rtol: relative tolerance for gradient equality.
delta: step size used for finite differences.
"""
if delta is None:
# Optimal stepsize for central difference is O(epsilon^{1/3}).
dtype = np_utils.result_type(*args)
epsilon = onp.finfo(dtype).eps
delta = epsilon ** (1.0 / 3.0)
theoretical, numerical = gradient_checker_v2.compute_gradient(
to_tf_fn(f), args, delta=delta)
self.assertAllClose(theoretical, numerical, check_dtypes=False, atol=atol,
rtol=rtol)
@contextmanager
def ignore_warning(**kw):
with warnings.catch_warnings():
warnings.filterwarnings("ignore", **kw)
yield
def disable(_):
def wrapper(self, *args, **kwargs):
self.skipTest('Test is disabled')
return wrapper