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 = npe.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(tf.nest.flatten(args),
tf.nest.flatten(new_args)):
if npe.most_precise_int_dtype(
old) != npe.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)
def _CompileAndCheck(self,
fun,
args_maker,
check_dtypes,
rtol=None,
atol=None,
check_eval_on_shapes=True,
check_incomplete_shape=False,
check_unknown_rank=True,
static_argnums=()):
"""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`.
"""
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
# `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)
python_ans = fun(*args)
python_shapes = tf.nest.map_structure(lambda x: onp.shape(x), python_ans)
onp_shapes = tf.nest.map_structure(lambda x: onp.shape(onp.asarray(x)),
python_ans)
self.assertEqual(python_shapes, onp_shapes)
cfun = npe.jit(wrapped_fun, static_argnums=static_argnums)
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(args, new_args):
if npe.most_precise_int_dtype(old) != npe.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)
if check_eval_on_shapes:
# Check that npe.eval_on_shapes can get complete output shapes given
# complete input shapes.
cfun = npe.eval_on_shapes(fun, static_argnums=static_argnums)
compiled_ans = cfun(*args)
flat_python_ans = tf.nest.flatten(python_ans)
flat_compiled_ans = tf.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(tf.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 = [tf.TensorSpec([None] * len(x.shape), x.dtype) for x in args]
cfun = npe.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 = [tf.TensorSpec(None, x.dtype) for x in args]
cfun = npe.jit(
fun, static_argnums=static_argnums, input_signature=specs)
compiled_ans = cfun(*args)
self.assertAllClose(python_ans, compiled_ans, check_dtypes, atol, rtol)
