text: # Copyright 2018 The TensorFlow Probability Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  """AutoGraph-based auto-batching frontend."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import contextlib
  import functools
  import gast
  import six
  
  import tensorflow.compat.v2 as tf
  
  from tensorflow_probability.python.experimental.auto_batching import allocation_strategy
  from tensorflow_probability.python.experimental.auto_batching import dsl
  from tensorflow_probability.python.experimental.auto_batching import instructions
  from tensorflow_probability.python.experimental.auto_batching import lowering
  from tensorflow_probability.python.experimental.auto_batching import stack_optimization as stack
  from tensorflow_probability.python.experimental.auto_batching import stackless as st
  from tensorflow_probability.python.experimental.auto_batching import tf_backend
  from tensorflow_probability.python.experimental.auto_batching import type_inference as ab_type_inference
  from tensorflow_probability.python.experimental.auto_batching import virtual_machine as vm
  
  # TODO(mdan): Move common converters under pyct; should have no other deps.
  # pylint:disable=g-direct-tensorflow-import
  from tensorflow.python.autograph.converters import break_statements
  from tensorflow.python.autograph.converters import continue_statements
  from tensorflow.python.autograph.converters import return_statements
  from tensorflow.python.autograph.core import converter
  from tensorflow.python.autograph.core import naming
  from tensorflow.python.autograph.pyct import anno
  from tensorflow.python.autograph.pyct import inspect_utils
  from tensorflow.python.autograph.pyct import parser
  from tensorflow.python.autograph.pyct import qual_names
  from tensorflow.python.autograph.pyct import templates
  from tensorflow.python.autograph.pyct import transformer
  from tensorflow.python.autograph.pyct.common_transformers import anf
  
  
  # For backward compatibility: compiler will be renamed to loader in TF 2.2+.
  try:
    from tensorflow.python.autograph.pyct import loader  # pylint:disable=g-import-not-at-top
  except ImportError:
    from tensorflow.python.autograph.pyct import compiler  # pylint:disable=g-import-not-at-top
  
    loader = compiler
    loader.load_ast = compiler.ast_to_object
  
  
  TF_BACKEND = tf_backend.TensorFlowBackend()
  
  
  def _parse_and_analyze(f, autobatch_functions):
    """Performs preliminary analyses and transformations.
  
    The goal is to massage the source program into a form on which
    the `_AutoBatchingTransformer` below will be successful.
  
    Args:
      f: Function to analyze
      autobatch_functions: List of Python `str` names of autobatched functions.
        Arguments to these functions will be canonicalized to variable references,
        but others will not.
  
    Returns:
      node: A Python AST node representing the function, suitable for
        passing to `_AutoBatchingTransformer.visit`
      entity_info: An AutoGraph `EntityInfo` object, with some information
        about `f`.  Required for initializing `_AutoBatchingTransformer`.
    """
    namespace = {}
  
    # Get the AST of the function
    future_features = inspect_utils.getfutureimports(f)
    node, _ = parser.parse_entity(f, future_features=future_features)
  
    # Boilerplate for AutoGraph transforms
    entity_info = transformer.EntityInfo(
        source_code='',
        source_file=None,
        future_features=future_features,
        namespace=namespace)
    program_ctx = converter.ProgramContext(
        options=converter.ConversionOptions(recursive=True),
        autograph_module=None)
    ctx = converter.EntityContext(
        namer=naming.Namer(namespace),
        entity_info=entity_info,
        program_ctx=program_ctx)
  
    # Canonicalize away break statements
    node = converter.standard_analysis(node, ctx, is_initial=True)
    node = break_statements.transform(node, ctx)
  
    # Canonicalize away continue statements
    node = converter.standard_analysis(node, ctx, is_initial=False)
    node = continue_statements.transform(node, ctx)
  
    # Force single returns
    node = converter.standard_analysis(node, ctx, is_initial=False)
    node = return_statements.transform(node, ctx, default_to_null_return=False)
  
    # Transform into ANF
    # Replacing if tests and autobatched function call arguments because
    # that's where divergence can happen.
    # Replacing all function calls because the downstream transformation
    # expects calls to lead directly to assignments.
    def maybe_replace_function_argument(parent, field_name, child):
      del field_name, child
      if not anno.hasanno(parent.func, anno.Basic.QN):
        return False
      func_name = anno.getanno(parent.func, anno.Basic.QN)
      if str(func_name) in autobatch_functions:
        return True
      return False
  
    anf_config = [
        (anf.ASTEdgePattern(gast.If, 'test', anf.ANY), anf.REPLACE),
        (anf.ASTEdgePattern(anf.ANY, anf.ANY, gast.Call), anf.REPLACE),
        (anf.ASTEdgePattern(gast.Call, 'args', anf.ANY),
         maybe_replace_function_argument),
    ]
    node = anf.transform(node, ctx, config=anf_config)
    node = converter.standard_analysis(node, ctx, is_initial=False)
  
    return node, ctx
  
  
  def _is_literal(node):
    """Detects whether the given node is a literal.
  
    This is surprisingly difficult to do robustly across versions of Python and
    gast, as the parsing of constants has changed, if I may, constantly.
  
    Args:
      node: The node whose status to check.
  
    Returns:
      literal: A Python `bool` giving whether the node is constant or not.
    """
    try:
      # TODO(b/140808434): Once we or TF decide to roll forward to gast 0.3, we
      # won't need this clause.
      # gast pre-0.3
      literals = (gast.Num, gast.Str, gast.Bytes, gast.Ellipsis,
                  gast.NameConstant)
      if isinstance(node, literals):
        return True
    except AttributeError:
      # gast 0.3+
      if isinstance(node, gast.Constant):
        return True
    # Python 2
    if isinstance(node, gast.Name) and node.id in ['True', 'False', 'None']:
      return True
    return False
  
  
  class _AutoBatchingTransformer(transformer.Base):
    """A subclass of `pyct.transformer.Base` implementing auto-batching.
  
    Specifically, converts some Python source P into a Python function that will
    invoke methods on a `dsl.ProgramBuilder` to register an auto-batchable version
    of P.  This class is private: use through the `transform` wrapper.
  
    The input is expected to be in A-normal form, so that all intermediate values
    are explicitly assigned.  As such, this class need only intercept function
    definitions, assignment statements, and control constructs.  At the moment,
    only supporting if, function call, and return.  TODO(axch): Expand the subset.
    """
  
    def __init__(self, known_functions, enclosing_names, ctx):
      """Initializes an _AutoBatchingTransformer.
  
      Args:
        known_functions: List of Python strings.  These are the names of
          auto-batched functions.  Calls to the functions on this list (when
          recognized) are transformed into `ProgramBuilder.call` methods.  Other
          function calls become `ProgramBuilder.primop`.
        enclosing_names: List of Python variables.  These are names to draw from
          the transformee's enclosing scope (without turning them into autobatch
          variables).
        ctx: An AutoGraph `autograph.converter.EntityContext` object describing
          the node to be transformed, as returned by `_parse_and_analyze`.
          Inherits from `autograph.transformer.Context`, see
          `autograph.transformer.Base`.
      """
      self.known_functions = known_functions
      self.enclosing_names = enclosing_names
      super(_AutoBatchingTransformer, self).__init__(ctx)
  
    def visit_FunctionDef(self, node):
      """Intercepts function definitions.
  
      Converts function definitions to the corresponding `ProgramBuilder.function`
      construction.
  
      Args:
        node: An `ast.AST` node representing the function to convert.
  
      Returns:
        node: An updated node, representing the result.
  
      Raises:
        ValueError: If the input node does not adhere to the restrictions,
          e.g., failing to have a `return` statement at the end.
      """
      # Check input form
      return_node = node.body[-1]
      if not isinstance(return_node, gast.Return):
        msg = 'Last node in function body should be Return, not {}.'
        raise ValueError(msg.format(return_node))
  
      # Convert all args to _tfp_autobatching_context_.param()
      local_declarations = []
      for arg in node.args.args:
        # print('Creating param declaration for', arg, arg.id, type(arg.id))
        local_declarations.append(templates.replace(
            'target = _tfp_autobatching_context_.param(name=target_name)',
            target=arg.id,
            target_name=gast.Str(arg.id))[0])
  
      # Visit the content of the function
      node = self.generic_visit(node)
  
      # Prepend the declarations
      node.body = local_declarations + node.body
  
      # Convert the function into a
      # `with _tfp_autobatching_context_.define_function()` block.
      # Wrap the `with` block into a function so additional information (namely,
      # the auto-batching `ProgramBuilder` and the `instruction.Function`s that
      # may be called in the body) can be passed in through regular Python
      # variable references.
      callable_function_names = [gast.Name(n, ctx=gast.Store(), annotation=None)
                                 for n in self.known_functions]
      node = templates.replace(
          '''
          def func(_tfp_autobatching_context_,
                   _tfp_autobatching_available_functions_):
            names = _tfp_autobatching_available_functions_
            with _tfp_autobatching_context_.define_function(func):
              body
            return func''',
          func=node.name,
          names=gast.List(callable_function_names, ctx=gast.Store()),
          body=node.body)[0]
  
      return node
  
    def visit_Assign(self, node):
      """Intercepts assignment statements.
  
      The input is expected to be in A-normal form, so every assignment represents
      at most one interesting computation, which is the AST node serving as the
      assignment's right-hand side.
  
      Args:
        node: An `ast.AST` node representing the assignment to convert.
  
      Returns:
        node: An updated node, representing the result.
  
      Raises:
        ValueError: If the node is a call to a function whose name cannot be
          determined.
      """
      # If we're assigning a constant value, use the `ProgramBuilder.const`
      # shorthand.
      if _is_literal(node.value):
        # Emit `_tfp_autobatching_context_.const`
        node = templates.replace(
            'target = _tfp_autobatching_context_.const(value)',
            target=self._assignment_construct(node.targets[0]),
            value=node.value)
        return node
  
      # If we're calling into a function, and the VM knows about the function, use
      # the `ProgramBuilder.call` instruction; otherwise, we'll emit a
      # `ProgramBuilder.primop` instruction, with a lambda staging the computation
      # to be performed.
      elif isinstance(node.value, gast.Call):
        if not anno.hasanno(node.value.func, anno.Basic.QN):
          # TODO(axch): Make calls to unidentified functions fall back to
          # primop, or flag them?
          raise ValueError('Cannot resolve name of the called function')
        func_name = anno.getanno(node.value.func, anno.Basic.QN)
        used_vars = node.value.args
        if str(func_name) in self.known_functions:
          # Emit `_tfp_autobatching_context_.call`
          node = templates.replace(
              'target = _tfp_autobatching_context_.call(func_name, vars_in=args)',
              target=self._assignment_construct(node.targets[0]),
              func_name=func_name,
              args=self._to_reference_list(used_vars))
          return node
        else:
          return self._primop(node, [func_name])
      else:
        return self._primop(node, [])
  
    def _primop(self, node, non_autobatch_names):
      scope = anno.getanno(node, anno.Static.SCOPE)
      # We want the root names of any attribute accesses; the accesses will happen
      # inside the expression.
      used_vars = set().union(*[qn.support_set for qn in scope.read])
      # Exclude names from the blacklist.
      used_vars = list(used_vars - set(non_autobatch_names))
      # Exclude names that will be available from the enclosing scope.
      used_vars = [v for v in used_vars if str(v) not in self.enclosing_names]
      # TODO(aqj) It should not be necessary to convert these QN objects to
      # strings, but empirically made the difference.
      stringified = [str(name) for name in used_vars]
      node = templates.replace(
          ('target = _tfp_autobatching_context_.primop('
           'lambda args: expr, vars_in=refs)'),
          target=self._assignment_construct(node.targets[0]),
          args=stringified,
          refs=self._to_reference_list(used_vars),
          expr=node.value)
      return node
  
    def _assignment_construct(self, target):
      if isinstance(target, (gast.Tuple, gast.List)):
        return templates.replace_as_expression(
            '_tfp_autobatching_context_(pat).pattern',
            pat=self._assignment_construct_recur(target))
      return self._assignment_construct_recur(target)
  
    def _assignment_construct_recur(self, target):
      if isinstance(target, (gast.Tuple, gast.List)):
        subs = [self._assignment_construct_recur(t) for t in target.elts]
        if isinstance(target, gast.Tuple):
          # Context is not Store anymore, because this section is constructing the
          # pattern object
          return gast.Tuple(subs, ctx=gast.Load())
        else:
          # Context is not Store anymore, because this section is constructing the
          # pattern object
          return gast.List(subs, ctx=gast.Load())
      return templates.replace_as_expression(
          '_tfp_autobatching_context_.var.name', name=target)
  
    def _to_reference(self, node):
      if isinstance(node, (gast.Name, qual_names.QN)):
        return templates.replace_as_expression(
            '_tfp_autobatching_context_.var.name', name=node)
      elif _is_literal(node):
        raise ValueError('TODO(axch): Support literals, not just variables')
      else:
        msg = 'Expected trivial node, got {}.  Is the input in A-normal form?'
        raise ValueError(msg.format(node))
  
    def _to_reference_list(self, names):
      return gast.List(
          [self._to_reference(name) for name in names], ctx=gast.Load())
  
    def visit_If(self, node):
      """Intercepts if statements.
  
      Converts each `if` to up to two separate `with` statements,
      `ProgramBuilder.if_(condition_variable)` and `ProgramBuilder.else_()`.  If
      the incoming `if` had one arm, returns the transformed AST node; if it had
      two, returns two nodes in a list.
  
      Args:
        node: An `ast.AST` node representing the `if` statement to convert.
  
      Returns:
        then_node: A node representing the `with`-guarded consequent branch.
        else_node: A node representing the `with`-guarded alternate branch,
          if present.
      """
      # Transform a branch
      # NOTE: this is a little hackery to make sure that prepending works
      # properly. Wrapping a list of statements in a Module ensures
      # that the AST-visiting machinery won't choke on, e.g., a list.
      then = self.generic_visit(gast.Module(node.body)).body
  
      # Construct header (goes in the `with`s).
      then_header = templates.replace_as_expression(
          '_tfp_autobatching_context_.if_(cond)',
          cond=self._to_reference(node.test))
  
      # Construct `with` node.
      # TODO(axch): Test that this form actually works with multiline bodies.
      then_node = templates.replace(
          'with header: body', header=then_header, body=then)[0]
  
      if node.orelse:
        orelse = self.generic_visit(gast.Module(node.orelse)).body
        orelse_header = templates.replace_as_expression(
            '_tfp_autobatching_context_.else_()')
        orelse_node = templates.replace(
            'with header: body', header=orelse_header, body=orelse)[0]
        # Return both
        return [then_node, orelse_node]
      else:
        return then_node
  
    def visit_Return(self, node):
      """Intercepts return statements.
  
      Args:
        node: An `ast.AST` node representing the `return` statement to convert.
  
      Returns:
        node: A node representing the result.
      """
      node = templates.replace_as_expression(
          '_tfp_autobatching_context_.return_(value)',
          value=self._to_reference(node.value))
      return gast.Expr(node)
  
  
  class Context(object):
    """Context object for auto-batching multiple Python functions together.
  
    Warning: This is alpha-grade software.  The exact subset of Python that can be
    successfully auto-batched is ill-specified and subject to change.
    Furthermore, the errors elicited by stepping outside this subset are likely to
    be confusing and misleading.  Use at your own risk.
  
    Usage:
    ```python
    ctx = frontend.Context()
  
    @ctx.batch(type_inference=lambda ...)
    def my_single_example_function_1(args):
      ...
  
    @ctx.batch(type_inference=lambda ...)
    def my_single_example_function_2(args):
      ...
  
    # etc
    ```
  
    Then calling any of the decorated functions will execute a batch computation.
    The decorated functions may call each other, including mutually recursively.
    See also the `batch` method.
  
    Limitations:
    - You must explicitly decorate every function to be auto-batched.
    - All calls to them must call them by name (no higher-order auto-batching).
    - Auto-batched functions must be defined with `def`, not `lambda`.
    """
  
    def __init__(self):
      """Initializes a `Context` object."""
      self._tagged_functions = []
      self._module = None
      self._in_dry_run = False
      # LRU-1 cache of compilation results, keyed by program, backend, and input
      # type signature.
      self._compile_cache = None
      # LRU-1 cache of lowering results, keyed by program, backend, and input
      # type signature.  These are separate because stackless mode requires
      # a non-lowered program, but lowering mutates (b/119122199).
      self._lowering_cache = None
  
    def _tag_function(self, function, type_inference):
      self._tagged_functions.append((function, type_inference))
      self._module = None
  
    @contextlib.contextmanager
    def _dry_run(self):
      old_dry_run = self._in_dry_run
      try:
        self._in_dry_run = True
        yield
      finally:
        self._in_dry_run = old_dry_run
  
    def batch(self, type_inference):
      """Decorates one function to auto-batch.
  
      The decorated function will run in batch.  It accepts all the same
      arguments, except:
      - All arguments must have an additional leading dimension for the batch.
        (By special dispensation, scalar inputs are promoted to shape [1], which
        then leads to broadcasting.)
      - All the arguments' sizes in the batch dimension must be the same, or 1.
        The latter are broadcast.
      - The returned value will also have a leading batch dimension, and
        will have the same size.
      - The batched function accepts an additional `bool` keyword argument
        `dry_run`.  If present and `True`, just calls the unbatched version,
        circumventing the auto-batching system.  This can be useful for
        debugging the program subject to auto-batching.
      - The batched function accepts an additional `bool` keyword argument
        `stackless`.  If present and `True`, invokes the stackless version of the
        auto-batching system.  This can be useful for avoiding stack maintenance
        overhead; but in general, it will recover less batching, and not work in
        graph-mode TensorFlow.
      - The batched function accepts an additional `int` keyword argument
        `max_stack_depth` specifying the maximum stack depth (default 15).
        Ignored in stackless execution.
      - The batched function accepts an additional keyword argument `backend`
        specifying the backend to use.  Must be an instance of
        `auto_batching.TensorFlowBackend` (default) or
        `auto_batching.NumpyBackend`.
      - The batched function accepts an additional keyword argument
        `block_code_cache`, a dict which allows the caching of basic block
        rewrites (i.e. `tf.function` + XLA) to live across calls to the
        autobatched function. The default value of `None` results in caching only
        within a given call to the batched function. Currently, stackless
        autobatching ignores the cache completely.
  
      Note: All functions called by the decorated function that need auto-batching
      must be decorated in the same `Context` before invoking any of them.
  
      Args:
        type_inference: A Python callable giving the type signature of the
          function being auto-batched.  The callable will be invoked with a single
          argument giving the list of `instructions.Type` objects describing the
          arguments at a particular call site, and must return a list of
          `instructions.Type` objects describing the values that call site will
          return.
  
      Returns:
        dec: A decorator that may be applied to a function with the given
          type signature to auto-batch it.
  
      Raises:
        ValueError: If the decorated function predictably cannot be auto-batched,
          e.g., name-clashing with another function already decorated in this
          `Context`.
      """
      return functools.partial(
          self.batch_uncurried, type_inference=type_inference)
  
    def batch_uncurried(self, function, type_inference):
      """A non-decorator version of `batch`, which see."""
      # TODO(axch): Friendly error message if `function` is not actually a
      # function.
      # TODO(axch): Friendly error message if `function` is a lambda and not a
      # `def`ed function.  Status quo will fail horribly on lambdas, because
      # the transformation requires the transformed functions be named.
      # Determine the name of the function being transformed
      name = function.__name__
      if name in self.function_names():
        msg = ('Cannot auto-batch the same-named function {} '
               'twice in the same `Context`.').format(name)
        raise ValueError(msg)
      self._tag_function(function, type_inference)
      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)
      return batched
  
    def function_names(self):
      return [f.__name__ for (f, _) in self._tagged_functions]
  
    def module(self):
      """Constructs an `instructions.Module` for this `Context`.
  
      Returns:
        module: An `instructions.Module` representing the batched computation
          defined by all the functions decorated with `batch` in this `Context` so
          far.
      """
      if self._module is not None:
        return self._module
      ab = dsl.ProgramBuilder()
      function_objects = []
      for function, type_inference in self._tagged_functions:
        declared = ab.declare_function(function.__name__, type_inference)
        function_objects.append(declared)
      for function, _ in self._tagged_functions:
        name = function.__name__
        node, ctx = _parse_and_analyze(function, self.function_names())
        node = _AutoBatchingTransformer(
            self.function_names(),
            [scoped_name for scoped_name, _ in _environment(function, [name])],
            ctx).visit(node)
        builder_module, _, _ = loader.load_ast(node)
        for scoped_name, val in _environment(function, [name]):
          builder_module.__dict__[scoped_name] = val
        builder = getattr(builder_module, name)
        builder(ab, function_objects)
      self._module = ab.module()
      return self._module
  
    def program(self, main):
      """Constructs an `instructions.Program` for this `Context`.
  
      This is a helper method, equivalent to `self.module().program(main)`.
  
      Args:
        main: Python string name of the function that should be the entry point.
  
      Returns:
        prog: An `instructions.Program` representing the batched computation
          defined by all the functions decorated with `batch` in this `Context` so
          far.  Suitable for downstream compilation with other passes in
          `auto_batching`.
  
      Raises:
        ValueError: If the intended `main` function was not decorated with
          `batch`.
      """
      return self.module().program(main)
  
    def program_compiled(self, main, sig=None, backend=None):
      """Constructs a compiled `instructions.Program` for this `Context`.
  
      This constructs the program with `self.program(main)`, and the performs type
      inference and optimization, to emit a result that can be executed by the
      stackless 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._compile_cache is not None:
        key, result = self._compile_cache
        if key == (module, main, sig, backend):
          return result
        else:
          # Clear the module cache as well, because of b/119122199
          self._module = 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
      result = allocation_strategy.optimize(typed)
      self._compile_cache = ((module, main, sig, backend), result)
      return result
  
    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 lowered_for_args(self, name, args, backend):
      """Helper for calling program_lowered that computes the type signature."""
      sig = ab_type_inference.signature(self.program(main=name), args, backend)
      lowered = self.program_lowered(main=name, sig=sig, backend=backend)
      return lowered
  
  
  def _environment(function, names_to_omit=()):
    """Yields the names and values visible from the function's scope."""
    str_names_to_omit = set(map(str, names_to_omit))
    for name, val in six.iteritems(six.get_function_globals(function)):
      if str(name) not in str_names_to_omit:
        yield name, val
    closure = six.get_function_closure(function)
    if closure is not None:
      freevars = six.get_function_code(function).co_freevars
      for name, cell in zip(freevars, closure):
        if str(name) not in str_names_to_omit:
          yield name, cell.cell_contents
  
  
  def _run_at_batch_size_one(backend, function, args):
    """Executes the given function, checking in- and out-batch shape."""
    for arg in args:
      for type_ in instructions.pattern_traverse(
          ab_type_inference.type_of_pattern(arg, backend),
          leaf_type=instructions.TensorType):
        shape = type_.shape
        if len(shape) >= 1 and shape[0] != 1:
          msg = 'Expecting input batch dimension of size 1; got {} of shape {}.'
          raise ValueError(msg.format(arg, shape))
    result = function(*args)
    for item in instructions.pattern_traverse(result):
      for type_ in instructions.pattern_traverse(
          ab_type_inference.type_of_pattern(item, backend),
          leaf_type=instructions.TensorType):
        shape = type_.shape
        if len(shape) >= 1 and shape[0] != 1:
          msg = 'Expecting result batch dimension of size 1; got {} of shape {}.'
          raise ValueError(msg.format(item, shape))
    return result
  
  
  def truthy(x):
    """Normalizes Tensor ranks for use in `if` conditions.
  
    This enables dry-runs of programs with control flow.  Usage: Program the
    conditions of `if` statements and `while` loops to have a batch dimension, and
    then wrap them with this function.  Example:
    ```python
    ctx = frontend.Context
    truthy = frontend.truthy
  
    @ctx.batch(type_inference=...)
    def my_abs(x):
      if truthy(x > 0):
        return x
      else:
        return -x
  
    my_abs([-5], dry_run=True)
    # returns [5] in Eager mode
    ```
  
    This is necessary because auto-batched programs still have a leading batch
    dimension (of size 1) even in dry-run mode, and a Tensor of shape [1] is not
    acceptable as the condition to an `if` or `while`.  However, the leading
    dimension is critical during batched execution; so conditions of ifs need to
    have rank 1 if running batched and rank 0 if running unbatched (i.e.,
    dry-run).  The `truthy` function arranges for this be happen (by detecting
    whether it is in dry-run mode or not).
  
    If you missed a spot where you should have used `truthy`, the error message
    will say `Non-scalar tensor <Tensor ...> cannot be converted to boolean.`
  
    Args:
      x: A Tensor.
  
    Returns:
      x: The Tensor `x` if we are in batch mode, or if the shape of `x` is
        anything other than `[1]`.  Otherwise returns the single scalar in `x` as
        a Tensor of scalar shape (which is acceptable in the conditions of `if`
        and `while` statements.
    """
    if ab_type_inference.is_inferring() or vm.is_staging() or st.is_running():
      return x
    elif isinstance(x, tf.Tensor):
      if x.shape == [1]:
        return x[0]
      else:
        return x
    else:
      return x
  