Exemplo n.º 1
0
 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)
Exemplo n.º 2
0
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)