<commit_before># 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 gast_util
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.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+.
# - `naming` will be moved to `core` in TF 2.2+.
try:
  from tensorflow.python.autograph.pyct import naming  # pylint:disable=g-import-not-at-top
except ImportError:
  from tensorflow.python.autograph.core import naming  # pylint:disable=g-import-not-at-top
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`.
  """
  # TODO(mdan): Replace all this boilerplate with FunctionTranspiler.
  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
  if hasattr(converter, 'EntityContext'):
    # TF 2.2-
    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)
  else:
    # TF 2.3+
    entity_info = transformer.EntityInfo(
        name=f.__name__,
        source_code='',
        source_file=None,
        future_features=future_features,
        namespace=namespace)
    program_ctx = converter.ProgramContext(
        options=converter.ConversionOptions(recursive=True),
        autograph_module=None)
    ctx = transformer.Context(
        info=entity_info,
        namer=naming.Namer(namespace),
        user_context=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


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<commit_msg>Internal cleanup: let each converter only perform the analysis it needs. Since the bottleneck is inside the conversion process, this is expected to speed it up noticeably.<commit_after># 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 gast_util
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.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
from tensorflow.python.autograph.pyct.static_analysis import activity


# For backward compatibility:
# - `compiler` will be renamed to `loader` in TF 2.2+.
# - `naming` will be moved to `core` in TF 2.2+.
# - 'converter.standard_analysis` will no longer be needed in TF 2.3+.
try:
  from tensorflow.python.autograph.pyct import naming  # pylint:disable=g-import-not-at-top
except ImportError:
  from tensorflow.python.autograph.core import naming  # pylint:disable=g-import-not-at-top
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

converter.standard_analysis = getattr(
    converter, 'standard_analysis', (lambda node, *_, **__: node))


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`.
  """
  # TODO(mdan): Replace all this boilerplate with FunctionTranspiler.
  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
  if hasattr(converter, 'EntityContext'):
    # TF 2.2-
    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)
  else:
    # TF 2.3+
    entity_info = transformer.EntityInfo(
        name=f.__name__,
        source_code='',
        source_file=None,
        future_features=future_features,
        namespace=namespace)
    program_ctx = converter.ProgramContext(
        options=converter.ConversionOptions(recursive=True),
        autograph_module=None)
    ctx = transformer.Context(
        info=entity_info,
        namer=naming.Namer(namespace),
        user_context=program_ctx)

  # Canonicalize away break statements
  node = converter.standard_analysis(node, ctx)
  node = break_statements.transform(node, ctx)

  # Canonicalize away continue statements
  node = converter.standard_analysis(node, ctx)
  node = continue_statements.transform(node, ctx)

  # Force single returns
  node = converter.standard_analysis(node, ctx)
  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 = qual_names.resolve(node)
  node = activity.resolve(node, ctx)

  return node, ctx


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
{'fraction_non_alphanumeric': 0.08072445019404915, 'fraction_numerical': 0.002328589909443726, 'mean_word_length': 4.671680117388115, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6541950', 'n_tokens_mistral': 4438, 'n_tokens_neox': 4242, 'n_words': 1307}