def _is_even_lowered_execute(inputs, backend):
prog = test_programs.is_even_function_calls()
alloc = allocation_strategy.optimize(prog)
lowered = lowering.lower_function_calls(alloc)
return list(vm.execute(
lowered, [inputs],
max_stack_depth=int(max(inputs)) + 3, backend=backend))
def _fibonacci_lowered_execute(inputs, backend):
prog = test_programs.fibonacci_function_calls()
alloc = allocation_strategy.optimize(prog)
lowered = lowering.lower_function_calls(alloc)
return list(vm.execute(
lowered, [inputs],
max_stack_depth=15, backend=backend))
def testProductTypeInferenceNumpy(self):
inputs = np.array([4, 5], dtype=np.int64)
outputs = np.array(([6, 7], [7, 8]), dtype=np.int64)
prog = test_programs.synthetic_pattern_variable_program(
include_types=False)
typed = type_inference.infer_types(prog, [inputs], NP_BACKEND)
expected_prog = test_programs.synthetic_pattern_variable_program()
self.assertSameTypes(expected_prog, typed)
alloc = allocation_strategy.optimize(typed)
lowered = lowering.lower_function_calls(alloc)
self.assertAllEqual(outputs, _execute(lowered, inputs, 15, NP_BACKEND))
def testAutoBatchingEvenOddNumpy(self):
for inputs, outputs in ([5], [False]), ([5, 6, 8, 9],
[False, True, True, False]):
inputs = np.array(inputs, dtype=np.int64)
outputs = np.array(outputs, dtype=np.bool_)
prog = even_odd_program()
# print(prog)
typed = type_inference.infer_types(prog, [inputs], NP_BACKEND)
# print(typed)
alloc = allocation_strategy.optimize(typed)
lowered = lowering.lower_function_calls(alloc)
# print(lowered)
self.assertAllEqual(outputs, _execute(lowered, inputs, 15, NP_BACKEND))
def testAutoBatchingFibonacciTF(self):
for inputs, outputs in ([5], [8]), ([5, 6, 8, 9], [8, 13, 34, 55]):
inputs = np.array(inputs, dtype=np.int32)
outputs = np.array(outputs, dtype=np.int32)
prog = fibonacci_program()
# print(prog)
inputs_t = tf.constant(inputs, dtype=np.int32)
typed = type_inference.infer_types(prog, [inputs_t], TF_BACKEND)
# print(typed)
alloc = allocation_strategy.optimize(typed)
lowered = lowering.lower_function_calls(alloc)
# print(lowered)
self.assertAllEqual(
outputs, self.evaluate(_execute(lowered, inputs_t, 15, TF_BACKEND)))
def testFibonacciTypeInferenceTF(self, dtype):
for inputs, outputs in ([5], [8]), ([5, 6, 8, 9], [8, 13, 34, 55]):
inputs = np.array(inputs, dtype=dtype)
outputs = np.array(outputs, dtype=dtype)
tf1.logging.debug('tf.fib {} {} {}'.format(
dtype, inputs.shape, outputs.shape))
inputs_t = tf.constant(inputs, dtype=dtype)
prog = test_programs.fibonacci_function_calls(include_types=False)
typed = type_inference.infer_types(prog, [inputs_t], TF_BACKEND)
expected_prog = test_programs.fibonacci_function_calls(dtype=dtype)
self.assertSameTypes(expected_prog, typed)
alloc = allocation_strategy.optimize(typed)
lowered = lowering.lower_function_calls(alloc)
self.assertAllEqual(
outputs, self.evaluate(_execute(lowered, inputs_t, 15, TF_BACKEND)))
def program_lowered(self, main, sig=None, backend=None):
"""Constructs a lowered `instructions.Program` for this `Context`.
This constructs the program with `self.program(main)`, and the performs type
inference, optimization, and lowering, to emit a result that can be executed
(or staged) by the auto-batching VM.
The point of having this as a method in its own right is that it caches the
compilation on the types of the arguments.
If either `sig` or `backend` are omitted or `None`, type inference is
skipped. The result is not executable, but it can be enlightening to
inspect.
Args:
main: Python string name of the function that should be the entry point.
sig: A `list` of (patterns of) `instructions.TensorType` aligned with
the formal parameters to `main`.
backend: Backend implementation.
Returns:
prog: An `instructions.Program` representing the batched computation
defined by all the functions decorated with `batch` in this `Context` so
far. Suitable for execution or staging on real data by the
auto-batching VM.
"""
module = self.module()
prog = module.program(main)
if self._lowering_cache is not None:
key, result = self._lowering_cache
if key == (module, main, sig, backend):
return result
else:
# Clear the module and compile caches as well, because of b/119122199
self._module = None
self._compile_cache = None
module = self.module()
prog = module.program(main)
if sig is not None and backend is not None:
typed = ab_type_inference.infer_types_from_signature(
prog, sig, backend)
else:
typed = prog
alloc = allocation_strategy.optimize(typed)
lowered = lowering.lower_function_calls(alloc)
result = stack.fuse_pop_push(lowered)
self._lowering_cache = ((module, main, sig, backend), result)
return result
def testAutoBatchingMultivalueTF(self):
input_ = np.array([1, 1, 1], dtype=np.int64)
output = ((np.array([1, 1, 1], dtype=np.int64),
np.array([3, 3, 3], dtype=np.int64)),
np.array([4, 4, 4], dtype=np.int64),
(np.array([5, 5, 5], dtype=np.int64),
np.array([6, 6, 6], dtype=np.int64)))
prog = synthetic_pattern_program()
# print(prog)
input_t = tf.constant(input_, dtype=np.int64)
typed = type_inference.infer_types(prog, [input_t], TF_BACKEND)
# print(typed)
alloc = allocation_strategy.optimize(typed)
lowered = lowering.lower_function_calls(alloc)
# print(lowered)
for expected, obtained in instructions.pattern_zip(
output, self.evaluate(_execute(lowered, input_t, 15, TF_BACKEND))):
self.assertAllEqual(expected, obtained)
def testIsEvenTypeInferenceTF(self, dtype):
for inputs, outputs in [([1], [False]),
([5, 6, 0, 3], [False, True, True, False])]:
inputs = np.array(inputs, dtype=dtype)
outputs = np.array(outputs, dtype=np.bool)
tf1.logging.debug('tf.even {} {} {}'.format(
dtype, inputs.shape, outputs.shape))
inputs_t = tf.constant(inputs, dtype=dtype)
prog = test_programs.is_even_function_calls(include_types=False)
typed = type_inference.infer_types(prog, [inputs_t], TF_BACKEND)
expected_prog = test_programs.is_even_function_calls(dtype=dtype)
self.assertSameTypes(expected_prog, typed)
alloc = allocation_strategy.optimize(typed)
lowered = lowering.lower_function_calls(alloc)
self.assertAllEqual(
outputs,
self.evaluate(_execute(
lowered, inputs_t, int(np.max(inputs)) + 3, TF_BACKEND)))
def testFibonacciTypeInferenceNumpy(self, dtype):
for inputs, outputs in ([5], [8]), ([5, 6, 8, 9], [8, 13, 34, 55]):
inputs = np.array(inputs, dtype=dtype)
outputs = np.array(outputs, dtype=dtype)
tf1.logging.debug('np.fib {} {} {}'.format(
dtype, inputs.shape, outputs.shape))
prog = test_programs.fibonacci_function_calls(include_types=False)
typed = type_inference.infer_types(prog, [inputs], NP_BACKEND)
expected_prog = test_programs.fibonacci_function_calls(dtype=dtype)
# We can only assert on the int64/float64 cases because numpy does
# not match-cast types on arithmetic with constants.
# i.e. (np.int32(0) - 1).dtype == np.int64
self.assertSameTypes(
expected_prog, typed, check_dtypes=dtype(0).nbytes == 8)
alloc = allocation_strategy.optimize(typed)
lowered = lowering.lower_function_calls(alloc)
self.assertAllEqual(
outputs, _execute(lowered, inputs, 15, NP_BACKEND))
def testIsEvenTypeInferenceNumpy(self, dtype):
for inputs, outputs in [([1], [False]),
([5, 6, 0, 3], [False, True, True, False])]:
inputs = np.array(inputs, dtype=dtype)
outputs = np.array(outputs, dtype=np.bool)
tf1.logging.debug('np.even {} {} {}'.format(
dtype, inputs.shape, outputs.shape))
prog = test_programs.is_even_function_calls(include_types=False)
typed = type_inference.infer_types(prog, [inputs], NP_BACKEND)
expected_prog = test_programs.is_even_function_calls(dtype=dtype)
# We can only assert on the int64/float64 cases because numpy does
# not match-cast types on arithmetic with constants.
# i.e. (np.int32(0) - 1).dtype == np.int64
self.assertSameTypes(
expected_prog, typed, check_dtypes=dtype(0).nbytes == 8)
alloc = allocation_strategy.optimize(typed)
lowered = lowering.lower_function_calls(alloc)
self.assertAllEqual(
outputs,
_execute(lowered, inputs, int(np.max(inputs)) + 3, NP_BACKEND))
def testAutoBatchingFibonacciNumpy(self):
for inputs, outputs in ([5], [8]), ([5, 6, 8, 9], [8, 13, 34, 55]):
# This test doesn't pass with int32 input types, because (apparently)
# numpy can't tell the difference between an ndarray of shape () and known
# dtype, and a scalar (literal) whose dtype needs to be inferred.
# To wit:
# (np.zeros((), dtype=np.int32) - 1).dtype == np.int64
# because that's somehow the best numpy can do, even though
# (np.zeros([6], dtype=np.int32) - 1).dtype == np.int32
# Needless to say, this messes up type inference for programs like
# Fibonacci whose unbatched input shape is scalar.
inputs = np.array(inputs, dtype=np.int64)
outputs = np.array(outputs, dtype=np.int64)
prog = fibonacci_program()
# print(prog)
typed = type_inference.infer_types(prog, [inputs], NP_BACKEND)
# print(typed)
alloc = allocation_strategy.optimize(typed)
lowered = lowering.lower_function_calls(alloc)
# print(lowered)
self.assertAllEqual(outputs, _execute(lowered, inputs, 15, NP_BACKEND))
def _is_even_stackless_execute(inputs, backend):
prog = test_programs.is_even_function_calls()
alloc = allocation_strategy.optimize(prog)
return stackless.execute(alloc, backend, None, inputs)
def assertAllocates(self, expected, prog):
allocated = allocation_strategy.optimize(prog)
self.assertEqual(expected, allocated.var_alloc)