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 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 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) tf.compat.v1.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 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) tf.compat.v1.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) tf.compat.v1.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 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 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) tf.compat.v1.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 assertAllocates(self, expected, prog): allocated = allocation_strategy.optimize(prog) self.assertEqual(expected, allocated.var_alloc)
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)