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 _constant_execute(inputs, backend): # Stack depth limit is 4 to accommodate the initial value and # two pushes to "answer". with tf.compat.v2.name_scope('constant_program'): return vm.execute(test_programs.constant_program(), [inputs], max_stack_depth=4, 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 batched(*args, **kwargs): """The batched function.""" # Accepting kwargs here because Python 2.7 gets confused otherwise. # See, e.g., # https://stackoverflow.com/questions/14003939/python-function-args-and-kwargs-with-other-specified-keyword-arguments # The whole strategy of adding arguments to the decorated function has # the unfortunate consequence that pylint uses the argument signature of # the decoratee to issue warnings, with the result that # disable=unexpected-keyword-arg needs to be added to use sites of # decorated functions, for example in `frontend_test.py`. # TODO(axch) Figure out how to avoid pylint warnings in callers. # Options: # - Add `backend` and `max_stack_depth` arguments to the `Context` # constructor, and use those as defaults here. # - Likely insufficient, because dynamic selection of the stack size # seems important. # - Carry the same information in a context manager, used like # with ctx.config(backend=TF_BACKEND, max_stack_depth=15): # some_decorated_function(normal, arguments) # - Just invoke batched functions from the `Context` object: # ctx.run(some_decorated_function, input_1, max_stack_depth=15) # - Maybe even make magic methods for the above? # ctx.some_decorated_function(input_1, max_stack_depth=15) max_stack_depth = kwargs.pop('max_stack_depth', 15) backend = kwargs.pop('backend', TF_BACKEND) dry_run = kwargs.pop('dry_run', False) stackless = kwargs.pop('stackless', False) block_code_cache = kwargs.pop('block_code_cache', {}) if self._in_dry_run: return _run_at_batch_size_one(backend, function, args) if dry_run: with self._dry_run(): return _run_at_batch_size_one(backend, function, args) if kwargs: msg = 'Auto-batched function given unexpected keyword arguments {}.' raise TypeError(msg.format(kwargs.keys())) sig = ab_type_inference.signature(self.program(main=name), args, backend) if stackless: compiled = self.program_compiled(main=name, sig=sig, backend=backend) return st.execute(compiled, backend, block_code_cache, *args) else: lowered = self.program_lowered(main=name, sig=sig, backend=backend) return vm.execute(lowered, args, max_stack_depth, backend, block_code_cache=block_code_cache)
def execute(batch_size, latent_size, data_size): data = np.random.normal(size=(data_size, latent_size)).astype(np.float32) def step_state(state): return state + np.sum(np.tensordot(data, state, ([1], [1]))) state = np.random.normal(size=(batch_size, latent_size)).astype(np.float32) def choose_depth(count): del count return 3 program = test_programs.pea_nuts_program((latent_size, ), choose_depth, step_state) input_counts = np.array([3] * batch_size) return vm.execute(program, [input_counts, state], 10, backend=NP_BACKEND)
def execute(_, backend): data = tf.random.normal(shape=(data_size, latent_size), dtype=np.float32) def step_state(state): return state + tf.reduce_sum( input_tensor=tf.tensordot(data, state, ([1], [1]))) state = tf.random.normal(shape=(batch_size, latent_size), dtype=np.float32) def choose_depth(count): del count return 2 program = test_programs.pea_nuts_program((latent_size, ), choose_depth, step_state) input_counts = np.array([3] * batch_size) return vm.execute(program, [input_counts, state], 10, backend=backend)
def _fibonacci_execute(inputs, backend): with tf.compat.v2.name_scope('fibonacci_program'): return vm.execute(test_programs.fibonacci_program(), [inputs], max_stack_depth=15, backend=backend)
def _product_type_execute(inputs, backend): with tf.compat.v2.name_scope('product_types_program'): return vm.execute(test_programs.synthetic_pattern_variable_program(), [inputs], max_stack_depth=4, backend=backend)
def _single_if_execute(inputs, backend): with tf.compat.v2.name_scope('single_if_program'): return vm.execute(test_programs.single_if_program(), [inputs], max_stack_depth=3, backend=backend)
def _execute(prog, inputs, stack_depth, backend): return vm.execute( prog, [inputs], max_stack_depth=stack_depth, backend=backend)