def testLoweringCache(self):
ctx = frontend.Context()
def my_type(args):
return args[0]
def function(x):
return x + x
ctx.batch_uncurried(function, my_type)
sig = [instructions.TensorType(np.int64, ())]
prog1 = ctx.program_lowered('function', sig, NP_BACKEND)
prog2 = ctx.program_lowered('function', sig, NP_BACKEND)
self.assertEqual(prog1, prog2)
sig2 = [instructions.TensorType(np.int32, ())]
prog3 = ctx.program_lowered('function', sig2, NP_BACKEND)
self.assertNotEqual(prog1, prog3)
def synthetic_pattern_variable_program(include_types=True):
"""A program that tests product types.
Args:
include_types: If False, we omit types on the variables, requiring a type
inference pass.
Returns:
program: `instructions.Program`.
"""
block = instructions.Block([
instructions.PrimOp(["inp"], "many", lambda x: (x + 1,
(x + 2, x + 3))),
instructions.PrimOp(["many"], ["one", "two"], lambda x: x),
], instructions.halt_op())
leaf = instructions.TensorType(np.int64, ())
the_vars = {
"inp": instructions.Type(leaf),
"many": instructions.Type((leaf, (leaf, leaf))),
"one": instructions.Type(leaf),
"two": instructions.Type((leaf, leaf)),
}
if not include_types:
_strip_types(the_vars)
return instructions.Program(instructions.ControlFlowGraph([block]), [],
the_vars, ["inp"], "two")
def _merge_tensor_type(old_type, obtained_type, backend):
"""Merges an updated instructions.TensorType for a single field."""
if old_type is None:
return obtained_type
# Update the inferred dtype.
if obtained_type is None:
return old_type
# One of old_type or obtained_type must be a TensorType because that's when
# the caller's pattern_map2 would call _merge_tensor_type.
if not isinstance(old_type, instructions.TensorType):
raise ValueError(
'Type mismatch: Expected struct type {}, got {}'.format(
old_type, obtained_type))
if not isinstance(obtained_type, instructions.TensorType):
raise ValueError(
'Type mismatch: Expected tensor type {}, got {}'.format(
old_type, obtained_type))
dtype = old_type.dtype
obtained_dtype = obtained_type.dtype
if dtype is None:
dtype = obtained_dtype
elif obtained_dtype is not None:
dtype = backend.merge_dtypes(dtype, obtained_dtype)
# Update the inferred shape.
shape = old_type.shape
obtained_shape = obtained_type.shape
if shape is None:
shape = obtained_shape
elif obtained_shape is not None:
shape = backend.merge_shapes(shape, obtained_shape)
return instructions.TensorType(dtype, shape)
def add_batch_dim_one_var(type_):
return instructions.Type(
instructions.pattern_map(
lambda t: instructions.TensorType(t.dtype,
(new_batch_dim, ) + t.shape),
type_.tensors,
leaf_type=instructions.TensorType))
def testFibonacciTF(self):
batch_fibo = frontend.Context().batch_uncurried(
fibonacci, lambda *args: instructions.TensorType(np.int64, ()))
input_2 = self._build_tensor(np.array([6, 7, 8], dtype=np.int64))
answer = batch_fibo(input_2, max_stack_depth=15, backend=TF_BACKEND)
self._check_batch_size(answer, 3)
self.assertAllEqual([8, 13, 21], self.evaluate(answer))
def var_init(max_stack_depth, initial_value):
type_ = inst.TensorType(initial_value.dtype, initial_value.shape[1:])
var = TF_BACKEND.create_variable(
None, inst.VariableAllocation.FULL, type_,
max_stack_depth, batch_size=initial_value.shape[0])
return var.update(
initial_value, TF_BACKEND.full_mask(initial_value.shape[0]))
def testSyncingAsExpected(self):
execution_counter_box = [0]
def count_executions(x):
execution_counter_box[0] += 1
return x
# Instrumented test program
def fibonacci_inst(n):
if n <= 1:
return 1
else:
left = fibonacci_inst(n - 2)
right = fibonacci_inst(n - 1)
return count_executions(left + right)
batch_fibo = frontend.Context().batch_uncurried(
fibonacci_inst,
lambda *args: instructions.TensorType(np.int64, ()))
self.assertEqual([3, 21, 5, 8],
list(
batch_fibo(np.array([3, 7, 4, 5], dtype=np.int64),
max_stack_depth=15,
backend=NP_BACKEND)))
# Expect 22 executions
# - 2 for type checking, plus
# - 20 because that's how many times 1 needs to be added to
# 1 to get 21.
self.assertEqual(22, execution_counter_box[0])
def testFibonacciNumpy(self):
batch_fibo = frontend.Context().batch_uncurried(
fibonacci,
lambda *args: instructions.TensorType(np.int64, ()))
self.assertEqual(
[13, 21, 34, 55],
list(batch_fibo(np.array([6, 7, 8, 9], dtype=np.int64),
max_stack_depth=15, backend=NP_BACKEND)))
def testFibonacciNumpyStackless(self):
if not tf.executing_eagerly():
return
batch_fibo = frontend.Context().batch_uncurried(
fibonacci,
lambda *args: instructions.TensorType(np.int64, ()))
self.assertEqual(
[3, 21, 5, 8],
list(batch_fibo(np.array([3, 7, 4, 5], dtype=np.int64),
max_stack_depth=15, backend=NP_BACKEND,
stackless=True)))
def type_of(self, t, preferred_dtype=None):
"""Returns the `instructions.Type` of `t`.
Args:
t: `np.ndarray` or a Python constant.
preferred_dtype: dtype to prefer, if `t` is a constant.
Returns:
vm_type: `instructions.TensorType` describing `t`
"""
t = np.array(t, dtype=preferred_dtype)
return instructions.TensorType(t.dtype, t.shape)
def type_of(self, t, dtype_hint=None):
"""Returns the `instructions.Type` of `t`.
Args:
t: `tf.Tensor` or a Python or numpy constant.
dtype_hint: dtype to prefer, if `t` is a constant.
Returns:
vm_type: `instructions.TensorType` describing `t`.
"""
if tf.executing_eagerly():
new_t = tf.convert_to_tensor(value=t, dtype=dtype_hint)
else:
with tf.Graph().as_default(): # Use a scratch graph.
new_t = tf.convert_to_tensor(value=t, dtype=dtype_hint)
dtype = new_t.dtype.base_dtype.as_numpy_dtype
shape = None if new_t.shape.ndims is None else tuple(new_t.shape.as_list())
return instructions.TensorType(dtype, shape)
def my_type(args): return instructions.TensorType(args[0].dtype, ())
def _strip_batch_dim(type_):
return instructions.pattern_map(
lambda t: instructions.TensorType(t.dtype, t.shape[1:]),
type_,
leaf_type=instructions.TensorType)
def pred_type(_): return instructions.TensorType(np.bool, ())
def pred_type(t):
return instructions.TensorType(np.bool, t[0].shape)
def testClosingOverTensorDoesntRaise(self):
x = tf.constant(0.)
def f(y):
return y * x
arg_types = [inst.Type([inst.TensorType(shape=[], dtype=np.float32)])]
TF_BACKEND.run_on_dummies(f, arg_types)
def pred_type(_): return instructions.TensorType(np.int32, ())
def my_type(_): int_ = instructions.TensorType(np.int64, ()) return ((int_, int_), int_, (int_, int_))
