# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# =============================================================================
# pylint: disable=g-classes-have-attributes
"""Contains a shim to allow using TF1 get_variable code in TF2."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import contextlib
import functools

from keras.engine import base_layer
from keras.utils import layer_utils
from keras.utils import tf_inspect
import tensorflow.compat.v2 as tf

from tensorflow.python.ops import variable_scope as vs  # pylint: disable=g-direct-tensorflow-import
from tensorflow.python.platform import tf_logging as logging  # pylint: disable=g-direct-tensorflow-import
from tensorflow.python.util.tf_export import keras_export  # pylint: disable=g-direct-tensorflow-import


def as_shape(shape):
  """Converts the given object to a TensorShape."""
  if isinstance(shape, tf.TensorShape):
    return shape
  else:
    return tf.TensorShape(shape)


def _is_callable_object(obj):
  return hasattr(obj, "__call__") and tf_inspect.ismethod(obj.__call__)


def _has_kwargs(fn):
  """Returns whether the passed callable has **kwargs in its signature.

  Args:
    fn: Function, or function-like object (e.g., result of `functools.partial`).

  Returns:
    `bool`: if `fn` has **kwargs in its signature.

  Raises:
     `TypeError`: If fn is not a Function, or function-like object.
  """
  if isinstance(fn, functools.partial):
    fn = fn.func
  elif _is_callable_object(fn):
    fn = fn.__call__
  elif not callable(fn):
    raise TypeError(
        "fn should be a function-like object, but is of type {}.".format(
            type(fn)))
  return tf_inspect.getfullargspec(fn).varkw is not None


def fn_args(fn):
  """Get argument names for function-like object.

  Args:
    fn: Function, or function-like object (e.g., result of `functools.partial`).

  Returns:
    `tuple` of string argument names.

  Raises:
    ValueError: if partial function has positionally bound arguments
  """
  if isinstance(fn, functools.partial):
    args = fn_args(fn.func)
    args = [a for a in args[len(fn.args):] if a not in (fn.keywords or [])]
  else:
    if hasattr(fn, "__call__") and tf_inspect.ismethod(fn.__call__):
      fn = fn.__call__
    args = tf_inspect.getfullargspec(fn).args
    if _is_bound_method(fn) and args:
      # If it's a bound method, it may or may not have a self/cls first
      # argument; for example, self could be captured in *args.
      # If it does have a positional argument, it is self/cls.
      args.pop(0)
  return tuple(args)


def _is_bound_method(fn):
  _, fn = tf.__internal__.decorator.unwrap(fn)
  return tf_inspect.ismethod(fn) and (fn.__self__ is not None)


def validate_synchronization_aggregation_trainable(
    synchronization, aggregation, trainable, name):
  """Given user-provided variable properties, sets defaults and validates."""
  if aggregation is None:
    aggregation = tf.compat.v1.VariableAggregation.NONE
  else:
    if not isinstance(aggregation,
                      (tf.compat.v1.VariableAggregation,
                       tf.VariableAggregation)):
      try:
        aggregation = tf.VariableAggregation(aggregation)
      except ValueError:
        raise ValueError(
            "Invalid variable aggregation mode: {} for variable: {}".format(
                aggregation, name))
  if synchronization is None:
    synchronization = tf.VariableSynchronization.AUTO
  else:
    try:
      synchronization = tf.VariableSynchronization(synchronization)
    except ValueError:
      raise ValueError(
          "Invalid variable synchronization mode: {} for variable: {}".format(
              synchronization, name))
  if trainable is None:
    trainable = synchronization != tf.VariableSynchronization.ON_READ
  return synchronization, aggregation, trainable


class _EagerVariableStore(tf.Module):
  """TF2-compatible VariableStore that avoids collections & tracks regularizers.

  New variable names and new variables can be created; all stored
  variables are initialized with the initializer passed to __init__.

  All variables get created in `tf.init_scope.` to avoid a bad
  interaction between `tf.function` `FuncGraph` internals, Keras
  Functional Models, and TPUStrategy variable initialization.

  Also, it always acts as if reuse is set to either "TRUE" or
  tf.compat.v1.AUTO_REUSE

  Attributes:
    vars: a dictionary with string names (same as passed in GetVar) as keys and
      the corresponding TensorFlow Variables as values.
    regularizers: a dictionary with string names as keys and the corresponding
      callables that return losses as values.
    layers: a dictionary with string names as keys and the corresponding
      nested keras layers as values.
  """

  def __init__(self):
    """Create a variable store."""
    self._vars = {}  # A dictionary of the stored TensorFlow variables.
    self._regularizers = {}  # A dict mapping var names to their regularizers.
    self._layers = {}  # A dictionary of stored keras layers.
    self._store_eager_variables = True

  @contextlib.contextmanager
  def scope(self):
    with vs.with_variable_store(self):
      yield

  def get_variable(
      self,
      name,
      shape=None,
      dtype=tf.float32,
      initializer=None,
      regularizer=None,
      reuse=None,
      trainable=None,
      collections=None,
      caching_device=None,
      partitioner=None,
      validate_shape=True,
      use_resource=None,
      custom_getter=None,
      constraint=None,
      synchronization=tf.VariableSynchronization.AUTO,
      aggregation=tf.compat.v1.VariableAggregation.NONE):
    """Gets an existing variable with these parameters or create a new one.

    If a variable with the given name is already stored, we return the stored
    variable. Otherwise, we create a new one.

    Set `reuse` to `True` when you only want to reuse existing Variables.
    Set `reuse` to None (the default) or tf.compat.v1.AUTO_REUSE when you want
    variables to be created if they don't exist or returned if they do.
    In this shim, `reuse` of `False` will be treated as auto-reuse.

    If initializer is `None` (the default), the default initializer passed in
    the constructor is used. If that one is `None` too, we use a new
    `glorot_uniform_initializer`. If initializer is a Tensor, we use
    it as a value and derive the shape from the initializer.

    If a partitioner is provided, a `PartitionedVariable` is returned.
    Accessing this object as a `Tensor` returns the shards concatenated along
    the partition axis.

    Some useful partitioners are available.  See, e.g.,
    `variable_axis_size_partitioner` and `min_max_variable_partitioner`.

    Args:
      name: The name of the new or existing variable.
      shape: Shape of the new or existing variable.
      dtype: Type of the new or existing variable (defaults to `DT_FLOAT`).
      initializer: Initializer for the variable.
      regularizer: A (Tensor -> Tensor or None) function; the result of applying
        it on a newly created variable will be added to the collection
        GraphKeys.REGULARIZATION_LOSSES and can be used for regularization.
      reuse: a Boolean, None, or tf.AUTO_REUSE. Controls reuse or creation of
        variables. When eager execution is enabled  this argument is always
        forced to be False.
      trainable: If `True` also add the variable to the graph collection
        `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). `trainable`
        defaults to `True`, unless `synchronization` is set to `ON_READ`, in
        which case it defaults to `False`.
      collections: List of graph collections keys to add the `Variable` to.
        Defaults to `[GraphKeys.GLOBAL_VARIABLES]` (see `tf.Variable`).
      caching_device: Optional device string or function describing where the
        Variable should be cached for reading.  Defaults to the Variable's
        device.  If not `None`, caches on another device.  Typical use is to
        cache on the device where the Ops using the `Variable` reside, to
        deduplicate copying through `Switch` and other conditional statements.
      partitioner: Optional callable that accepts a fully defined `TensorShape`
        and dtype of the `Variable` to be created, and returns a list of
        partitions for each axis (currently only one axis can be partitioned).
      validate_shape: If False, allows the variable to be initialized with a
        value of unknown shape. If True, the default, the shape of initial_value
        must be known.
      use_resource: If False, creates a regular Variable. If True, creates
        instead an experimental ResourceVariable which has well-defined
        semantics. Defaults to False (will later change to True). When eager
        execution is enabled this argument is always forced to be true.
      custom_getter: Callable that takes as a first argument the true getter,
        and allows overwriting the internal get_variable method. The signature
        of `custom_getter` should match that of this method,
        but the most future-proof version will allow for changes: `def
          custom_getter(getter, *args, **kwargs)`.  Direct access to
        all `get_variable` parameters is also allowed: `def
          custom_getter(getter, name, *args, **kwargs)`.  A simple identity
        custom getter that simply creates variables with modified names is:
          ```python
        def custom_getter(getter, name, *args, **kwargs): return getter(name +
          '_suffix', *args, **kwargs) ```
      constraint: An optional projection function to be applied to the variable
        after being updated by an `Optimizer` (e.g. used to implement norm
        constraints or value constraints for layer weights). The function must
        take as input the unprojected Tensor representing the value of the
        variable and return the Tensor for the projected value (which must have
        the same shape). Constraints are not safe to use when doing asynchronous
        distributed training.
      synchronization: Indicates when a distributed a variable will be
        aggregated. Accepted values are constants defined in the class
        `tf.VariableSynchronization`. By default the synchronization is set to
        `AUTO` and the current `DistributionStrategy` chooses when to
        synchronize.
      aggregation: Indicates how a distributed variable will be aggregated.
        Accepted values are constants defined in the class
        `tf.VariableAggregation`.

    Returns:
      The created or existing `Variable` (or `PartitionedVariable`, if a
      partitioner was used).

    Raises:
      ValueError: when creating a new variable and shape is not declared,
        when reusing a variable and specifying a conflicting shape,
        or when violating reuse during variable creation.
      RuntimeError: when eager execution is enabled and not called from an
        EagerVariableStore.
    """
    if custom_getter is not None and not callable(custom_getter):
      raise ValueError("Passed a custom_getter which is not callable: %s" %
                       custom_getter)

    with tf.init_scope():
      if tf.executing_eagerly():
        # Variable creation and initialization takes place in `init_scope`s;
        # as such, if an `init_scope` lifts us into the eager context, then we
        # need to use `ResourceVariable`s.
        use_resource = True

    # Note that it's fine to reuse eager variables whose initialization was
    # lifted from a function-building graph into the eager context (that's why
    # the following clause is not wrapped in an `init_scope`); lifted variables
    # are tracked by the graph's `VariableStore`.
    if not reuse:
      reuse = tf.compat.v1.AUTO_REUSE

    # If a *_ref type is passed in an error would be triggered further down the
    # stack. We prevent this using base_dtype to get a non-ref version of the
    # type, before doing anything else. When _ref types are removed in favor of
    # resources, this line can be removed.
    try:
      dtype = dtype.base_dtype
    except AttributeError:
      # .base_dtype not existing means that we will try and use the raw dtype
      # which was passed in - this might be a NumPy type which is valid.
      pass

    # This is the main logic of get_variable.  However, custom_getter
    # may override this logic.  So we save it as a callable and pass
    # it to custom_getter.
    # Note: the parameters of _true_getter, and their documentation, match
    # *exactly* item-for-item with the docstring of this method.
    def _true_getter(  # pylint: disable=missing-docstring
        name,
        shape=None,
        dtype=tf.float32,
        initializer=None,
        regularizer=None,
        reuse=None,
        trainable=None,
        collections=None,  # pylint: disable=unused-argument
        caching_device=None,
        partitioner=None,
        validate_shape=True,
        use_resource=None,  # pylint: disable=unused-argument
        constraint=None,
        synchronization=tf.VariableSynchronization.AUTO,
        aggregation=tf.compat.v1.VariableAggregation.NONE):
      # Partitioned variable currently unsupported w/ the shim
      if partitioner is not None:
        raise ValueError(
            "`partitioner` arg for `get_variable` is unsupported in TF2."
            "File a bug if you need help. You passed %s" % partitioner)

      # Single variable case
      if "%s/part_0" % name in self._vars:
        raise ValueError(
            "No partitioner was provided, but a partitioned version of the "
            "variable was found: %s/part_0. Perhaps a variable of the same "
            "name was already created with partitioning?" % name)

      return self._get_single_variable(
          name=name,
          shape=shape,
          dtype=dtype,
          initializer=initializer,
          regularizer=regularizer,
          reuse=reuse,
          trainable=trainable,
          caching_device=caching_device,
          validate_shape=validate_shape,
          constraint=constraint,
          synchronization=synchronization,
          aggregation=aggregation)

    synchronization, aggregation, trainable = (
        validate_synchronization_aggregation_trainable(
            synchronization, aggregation, trainable, name))

    if custom_getter is not None:
      # Handle backwards compatibility with getter arguments that were added
      # to the API after users started writing custom getters.
      custom_getter_kwargs = {
          "getter": _true_getter,
          "name": name,
          "shape": shape,
          "dtype": dtype,
          "initializer": initializer,
          "regularizer": regularizer,
          "reuse": reuse,
          "trainable": trainable,
          "collections": collections,
          "caching_device": caching_device,
          "partitioner": partitioner,
          "validate_shape": validate_shape,
          "use_resource": use_resource,
          "synchronization": synchronization,
          "aggregation": aggregation,
      }
      # `fn_args` and `has_kwargs` can handle functions, `functools.partial`,
      # `lambda`.
      if ("constraint" in fn_args(custom_getter) or
          _has_kwargs(custom_getter)):
        custom_getter_kwargs["constraint"] = constraint
      return custom_getter(**custom_getter_kwargs)
    else:
      return _true_getter(
          name,
          shape=shape,
          dtype=dtype,
          initializer=initializer,
          regularizer=regularizer,
          reuse=reuse,
          trainable=trainable,
          collections=collections,
          caching_device=caching_device,
          partitioner=partitioner,
          validate_shape=validate_shape,
          use_resource=use_resource,
          constraint=constraint,
          synchronization=synchronization,
          aggregation=aggregation)

  def _get_single_variable(
      self,
      name,
      shape=None,
      dtype=tf.float32,
      initializer=None,
      regularizer=None,
      partition_info=None,
      reuse=None,
      trainable=None,
      caching_device=None,
      validate_shape=True,
      constraint=None,
      synchronization=tf.VariableSynchronization.AUTO,
      aggregation=tf.compat.v1.VariableAggregation.NONE):
    """Get or create a single Variable (e.g.

    a shard or entire variable).

    See the documentation of get_variable above (ignore partitioning components)
    for details.

    Args:
      name: see get_variable.
      shape: see get_variable.
      dtype: see get_variable.
      initializer: see get_variable.
      regularizer: see get_variable.
      partition_info: _PartitionInfo object.
      reuse: see get_variable.
      trainable: see get_variable.
      caching_device: see get_variable.
      validate_shape: see get_variable.
      constraint: see get_variable.
      synchronization: see get_variable.
      aggregation: see get_variable.

    Returns:
      A Variable.  See documentation of get_variable above.

    Raises:
      ValueError: See documentation of get_variable above.
    """
    # Set to true if initializer is a constant.
    initializing_from_value = False
    if initializer is not None and not callable(initializer):
      initializing_from_value = True
    if shape is not None and initializing_from_value:
      raise ValueError("If initializer is a constant, do not specify shape.")

    dtype = tf.as_dtype(dtype)
    shape = as_shape(shape)

    if name in self._vars:
      # Here we handle the case when returning an existing variable.
      found_var = self._vars[name]
      if not shape.is_compatible_with(found_var.get_shape()):
        raise ValueError("Trying to share variable %s, but specified shape %s"
                         " and found shape %s." %
                         (name, shape, found_var.get_shape()))
      if not dtype.is_compatible_with(found_var.dtype):
        dtype_str = dtype.name
        found_type_str = found_var.dtype.name
        raise ValueError("Trying to share variable %s, but specified dtype %s"
                         " and found dtype %s." %
                         (name, dtype_str, found_type_str))
      return found_var

    # The code below handles only the case of creating a new variable.
    if reuse is True:  # pylint: disable=g-bool-id-comparison
      raise ValueError("Variable %s does not exist, or was not created with "
                       "tf.get_variable(). Did you mean to set "
                       "reuse=tf.AUTO_REUSE in VarScope?" % name)

    # Create the tensor to initialize the variable with default value.
    if initializer is None:
      initializer, initializing_from_value = self._get_default_initializer(
          name=name, shape=shape, dtype=dtype)
    # Enter an init scope when creating the initializer.
    with tf.init_scope():
      if initializing_from_value:
        init_val = initializer
        variable_dtype = None
      else:
        # Instantiate initializer if provided initializer is a type object.
        if tf_inspect.isclass(initializer):
          initializer = initializer()
        if shape.is_fully_defined():
          if "partition_info" in tf_inspect.getargspec(initializer).args:
            init_val = functools.partial(initializer,
                                         shape.as_list(),
                                         dtype=dtype,
                                         partition_info=partition_info)
          else:
            init_val = functools.partial(initializer,
                                         shape.as_list(), dtype=dtype)
          variable_dtype = dtype.base_dtype
        else:
          init_val = initializer
          variable_dtype = None

    # Create the variable (Always eagerly as a workaround for a strange
    # tpu / funcgraph / keras functional model interaction )
    with tf.init_scope():
      v = tf.Variable(
          initial_value=init_val,
          name=name,
          trainable=trainable,
          caching_device=caching_device,
          dtype=variable_dtype,
          validate_shape=validate_shape,
          constraint=constraint,
          synchronization=synchronization,
          aggregation=aggregation)

    self._vars[name] = v
    logging.vlog(1, "Created variable %s with shape %s and init %s", v.name,
                 format(shape), initializer)

    # Run the regularizer if requested and save the resulting loss.
    if regularizer:
      self.add_regularizer(v, regularizer)

    return v

  def get_or_create_layer(self, name, create_layer_method):
    if name not in self._layers:
      layer = create_layer_method()
      self._layers[name] = layer
      if isinstance(layer, base_layer.Layer):
        self._regularizers[name] = lambda: tf.math.reduce_sum(layer.losses)
    return self._layers[name]

  def add_regularizer(self, var, regularizer):
    self._regularizers[var.name] = functools.partial(regularizer, var)

  # Initialize variable when no initializer provided
  def _get_default_initializer(self, name, shape=None, dtype=tf.float32):
    """Provide a default initializer and a corresponding value.

    Args:
      name: see get_variable.
      shape: see get_variable.
      dtype: see get_variable.

    Returns:
      initializer and initializing_from_value. See get_variable above.

    Raises:
      ValueError: When giving unsupported dtype.
    """
    del shape
    # If dtype is DT_FLOAT, provide a uniform unit scaling initializer
    if dtype.is_floating:
      initializer = tf.compat.v1.glorot_uniform_initializer()
      initializing_from_value = False
    # If dtype is DT_INT/DT_UINT, provide a default value `zero`
    # If dtype is DT_BOOL, provide a default value `FALSE`
    elif (dtype.is_integer or dtype.is_unsigned or dtype.is_bool or
          dtype == tf.string):
      initializer = tf.compat.v1.zeros_initializer()
      initializing_from_value = False
    # NOTES:Do we need to support for handling DT_STRING and DT_COMPLEX here?
    else:
      raise ValueError("An initializer for variable %s of %s is required" %
                       (name, dtype.base_dtype))

    return initializer, initializing_from_value


@keras_export(v1=["keras.utils.track_tf1_style_variables"])
def track_tf1_style_variables(method):
  """Wrap layer & module methods in this decorator to capture tf1-style weights.

  Decorating a `tf.keras.Layer`'s  or `tf.Module`'s methods with this
  decorator will cause the layer/module to track weights created/used
  via `tf.compat.v1.get_variable` (and by extension `tf.compat.v1.layers`)
  inside the decorated method.

  In addition to tracking the weights themselves under the standard
  `layer.variable`/`module.variable`/etc. properties, if the method belongs
  to a `tf.keras.Layer` then any regularization losses specified via the
  `get_variable` or `tf.compat.v1.layers` regularizer arguments will get
  tracked by the layer under the standard `layer.losses` property.

  This tracking enables using large classes of TF1-style model-forward-pass
  code inside of Keras layers or `tf.Modules` in TF2 with TF2 behaviors enabled.

  Example of capturing tf.compat.v1.layer-based modeling code as a Keras layer:

  ```python
  class WrappedDoubleDenseLayer(tf.keras.layers.Layer):

    def __init__(self, units, *args, **kwargs):
      super().__init__(*args, **kwargs)
      self.units = units

    @tf.compat.v1.keras.utils.track_tf1_style_variables
    def call(self, inputs):
      with tf.compat.v1.variable_scope("double_dense_layer"):
        out = tf.compat.v1.layers.dense(
            inputs, self.units, name="dense_one",
            kernel_initializer=tf.compat.v1.random_normal_initializer,
            kernel_regularizer="l2")
        out = tf.compat.v1.layers.dense(
            out, self.units, name="dense_two",
            kernel_initializer=tf.compat.v1.random_normal_initializer(),
            kernel_regularizer="l2")
      return out

  # Create a layer that can be used as a standard keras layer
  layer = WrappedDoubleDenseLayer(10)

  # call the layer on inputs
  layer(...)

  # Variables created/used within the scope will be tracked by the layer
  layer.weights
  layer.trainable_variables

  # Regularization losses will be captured in layer.losses after a call,
  # just like any other Keras layer
  reg_losses = layer.losses
  ```

  Example of capturing tf.compat.v1.get_variable-based modeling code as
  a Keras layer:

  ```python
  class WrappedDoubleDenseLayer(tf.keras.layers.Layer):

    def __init__(self, units, *args, **kwargs):
      super().__init__(*args, **kwargs)
      self.units = units

    @tf.compat.v1.keras.utils.track_tf1_style_variables
    def call(self, inputs):
      out = inputs
      with tf.compat.v1.variable_scope("double_dense_layer"):
        with tf.compat.v1.variable_scope("dense_one"):
          # The weights are created with a `regularizer`,
          # so the layer should track their regularization losses
          kernel = tf.compat.v1.get_variable(
              shape=[out.shape[-1], self.units],
              regularizer=regularizers.L2(),
              initializer=init_ops.ones_initializer(),
              name="kernel")
          bias = tf.compat.v1.get_variable(
              shape=[self.units,],
              initializer=init_ops.zeros_initializer(),
              name="bias")
          out = tf.compat.v1.math.matmul(out, kernel)
          out = tf.compat.v1.nn.bias_add(out, bias)
        with tf.compat.v1.variable_scope("dense_two"):
          kernel = tf.compat.v1.get_variable(
              shape=[out.shape[-1], self.units],
              regularizer=regularizers.L2(),
              initializer=init_ops.ones_initializer(),
              name="kernel")
          bias = tf.compat.v1.get_variable(
              shape=[self.units,],
              initializer=init_ops.zeros_initializer(),
              name="bias")
          out = tf.compat.v1.math.matmul(out, kernel)
          out = tf.compat.v1.nn.bias_add(out, bias)
      return out

  # Create a layer that can be used as a standard keras layer
  layer = WrappedDoubleDenseLayer(10)

  # call the layer on inputs
  layer(...)

  # Variables created/used within the scope will be tracked by the layer
  layer.weights
  layer.trainable_variables

  # Regularization losses will be captured in layer.losses after a call,
  # just like any other Keras layer
  reg_losses = layer.losses
  ```

  Regularization losses:
    Any regularizers specified in the `get_variable` calls or `compat.v1.layer`
    creations will get captured if they occur in your decorated method
    and the method belongs to a `tf.keras.Layer`/`tf.keras.Module`.
    Regularization losses
    are accessible in `layer.losses` after a call just like in a standard
    Keras layer, and will be captured by any model that includes this layer.
    Regularization losses attached to Keras layers/models set as attributes
    of your layer will also get captured in the standard Keras regularization
    loss tracking.

    (While Modules have no `losses` property, no-arg callables to compute
     the regularization losses may be tracked as dict values in a private
     `module._tf1_style_var_store._regularizers` property, but only for
     `tf.compat.v1.layers` and `get_variable` weights and not for any other
     nested Keras layers/tf.Modules)

  Variable scope / variable reuse:
    variable-scope based reuse in your decorated method will be respected,
    and work like variable-scope based reuse in TF1.

  Variable Names/Pre-trained checkpoint loading:
    Variable naming from get_variable and `compat.v1.layer` layers will match
    the TF1 names, so you should be able to re-use your old name-based
    checkpoints. Variable naming for Keras layers/models or for variables
    created by `tf.Variable` may change when going to eager execution.

  Training Arg if you decorate `layer.call`:
    Keras will pass a `training` arg to this layer if `call` contains
    a `training` arg or a `**kwargs` varargs in its call signature,
    similarly to how keras passes `training` to other layers in TF2 that have
    similar signatures in their `call` implementations.
    See more details in the docs
    on `tf.keras.layers.Layer` to understand what will be passed and when.
    Note: tf.compat.v1.layers are usually not called with `training=None`,
    so the training arg to `forward_pass` might not feed through to them
    unless you pass it to their calls explicitly.

  Caveats:
    * TF2 will not prune unused variable updates (or unused outputs). You may
      need to adjust your forward pass code to avoid computations or variable
      updates that you don't intend to use.
    * Avoid Nesting variable creation in tf.function inside of
      methods decorated with `track_tf1_style_variables`
      While the method may safely be used from inside a `tf.function`, using
      a function inside of a decorated method may break the variable scoping.
    * This decorator only adds implicit tracking for legacy tf1-style
      get_variable / compat.v1.layers usage.
      If you would like to use nested Keras layers/models
      inside the decorated method, you need to
      assign them as attributes of your layer so that Keras/Module's standard
      object-oriented weights (and loss tracking for layers) will kick in.
      See the intro to modules, layers, and models
      [guide](https://www.tensorflow.org/guide/intro_to_modules) for more info.
      As a backup, the `compat.v1.keras.utils.get_or_create_layer` method will
      ease tracking nested keras model weights and losses for existing TF1 code,
      but new code should use explicit tracking.

  Args:
    method: The method to decorate. This should belong to a custom tf.Module,
    tf.keras.layers.Layer, or tf.keras.Model.

  Returns:
    The decorated method.
  """

  def _method_wrapper(self, *args, **kwargs):
    var_store = getattr(self, "_tf1_style_var_store", None)
    if not var_store:
      if not isinstance(self, tf.Module):
        # Raise an error if you incorrectly decorate a method
        # that is not a method of a Module, Layer, or Model:
        raise ValueError(
            "`@tf.compat.v1.keras.utils.track_tf1_layers_and_variables` must "
            "be applied to a method of a subclassed `tf.Module`, "
            "`tf.keras.layers.Layer`, or `tf.keras.Model` and which takes "
            "`self` as the first argument. But, the first argument passed "
            "to the decorated method was {}, which does not "
            "extend Module, Layer, or Model.".format(self))
      var_store = _EagerVariableStore()
      self._tf1_style_var_store = var_store  # pylint: disable=protected-access

    existing_regularized_variables = set(var_store._regularizers.keys())  # pylint: disable=protected-access
    with var_store.scope():
      out = method(self, *args, **kwargs)

    # If this is a layer method, add the regularization losses
    # to the layer for any newly-created regularized variables
    if isinstance(self, base_layer.Layer):
      for var_name, regularizer in var_store._regularizers.items():  # pylint: disable=protected-access
        if var_name not in existing_regularized_variables:
          self.add_loss(regularizer)

    return out

  return tf.__internal__.decorator.make_decorator(
      target=method, decorator_func=_method_wrapper)


class VariableScopeLayer(base_layer.Layer):
  """Wrapper Layer to capture `compat.v1.get_variable` and `compat.v1.layers`.

  This shim layer allows using large sets of TF1 model-forward-pass code as a
  Keras layer that works in TF2 with TF2 behaviors enabled. It will capture
  both weights and regularization losses of your forward-pass code. To use it,
  override this class and put your TF1 model's forward pass inside your
  implementation for `forward_pass`. (Unlike standard custom Keras layers,
  do not override `call`.)

  Below are some examples, and then more details on the functionality of this
  shim layer to wrap TF1 model forward passes.

  Example of capturing tf.compat.v1.layer-based modeling code as a Keras layer:

  ```python
  class WrappedDoubleDenseLayer(variable_scope_shim.VariableScopeLayer):

    def __init__(self, units, *args, **kwargs):
      super().__init__(*args, **kwargs)
      self.units = units

    def forward_pass(self, inputs):
      with variable_scope.variable_scope("double_dense_layer"):
        out = tf.compat.v1.layers.dense(
            inputs, self.units, name="dense_one",
            kernel_initializer=tf.compat.v1.random_normal_initializer,
            kernel_regularizer="l2")
        out = tf.compat.v1.layers.dense(
            out, self.units, name="dense_two",
            kernel_initializer=tf.compat.v1.random_normal_initializer(),
            kernel_regularizer="l2")
      return out

  # Create a layer that can be used as a standard keras layer
  layer = WrappedDoubleDenseLayer(10)

  # call the layer on inputs
  layer(...)

  # Variables created/used within the scope will be tracked by the layer
  layer.weights
  layer.trainable_variables

  # Regularization losses will be captured in layer.losses after a call,
  # just like any other Keras layer
  reg_losses = layer.losses
  ```

  Example of capturing tf.compat.v1.get_variable-based modeling code as
  a Keras layer:

  ```python
  class WrappedDoubleDenseLayer(variable_scope_shim.VariableScopeLayer):

    def __init__(self, units, *args, **kwargs):
      super().__init__(*args, **kwargs)
      self.units = units

    def forward_pass(self, inputs):
      out = inputs
      with tf.compat.v1.variable_scope("double_dense_layer"):
        with tf.compat.v1.variable_scope("dense_one"):
          # The weights are created with a `regularizer`,
          # so the layer should track their regularization losses
          kernel = tf.compat.v1.get_variable(
              shape=[out.shape[-1], self.units],
              regularizer=regularizers.L2(),
              initializer=init_ops.ones_initializer(),
              name="kernel")
          bias = tf.compat.v1.get_variable(
              shape=[self.units,],
              initializer=init_ops.zeros_initializer(),
              name="bias")
          out = tf.compat.v1.math.matmul(out, kernel)
          out = tf.compat.v1.nn.bias_add(out, bias)
        with tf.compat.v1.variable_scope("dense_two"):
          kernel = tf.compat.v1.get_variable(
              shape=[out.shape[-1], self.units],
              regularizer=regularizers.L2(),
              initializer=init_ops.ones_initializer(),
              name="kernel")
          bias = tf.compat.v1.get_variable(
              shape=[self.units,],
              initializer=init_ops.zeros_initializer(),
              name="bias")
          out = tf.compat.v1.math.matmul(out, kernel)
          out = tf.compat.v1.nn.bias_add(out, bias)
      return out

  # Create a layer that can be used as a standard keras layer
  layer = WrappedDoubleDenseLayer(10)

  # call the layer on inputs
  layer(...)

  # Variables created/used within the scope will be tracked by the layer
  layer.weights
  layer.trainable_variables

  # Regularization losses will be captured in layer.losses after a call,
  # just like any other Keras layer
  reg_losses = layer.losses
  ```

  Regularization losses:
    Any regularizers specified in the `get_variable` calls or `compat.v1.layer`
    creations will get captured by this wrapper layer. Regularization losses
    are accessible in `layer.losses` after a call just like in a standard
    Keras layer, and will be captured by any model that includes this layer.
    Regularization losses attached to Keras layers/models set as attributes
    of your layer will also get captured in the standard Keras regularization
    loss tracking.

  Variable scope / variable reuse:
    variable-scope based reuse in the `forward_pass` will be respected,
    and work like variable-scope based reuse in TF1.

  Variable Names/Pre-trained checkpoint loading:
    Variable naming from get_variable and `compat.v1.layer` layers will match
    the TF1 names, so you should be able to re-use your old name-based
    checkpoints. Variable naming for Keras layers/models or for variables
    created by `tf.Variable` may change when going to eager execution.

  Training Arg in `forward_pass`:
    Keras will pass a `training` arg to this layer if `forward_pass` contains
    a `training` arg or a `**kwargs` varargs in its call signature,
    similarly to how keras passes `training` to other layers in TF2 that have
    similar signatures in their `call` implementations.
    See more details in the docs
    on `tf.keras.layers.Layer` to understand what will be passed and when.
    Note: tf.compat.v1.layers are usually not called with `training=None`,
    so the training arg to `forward_pass` might not feed through to them
    unless you pass it to their calls explicitly.

  Call signature of the forward pass:
    The semantics of the forward pass signature match the standard
    Keras layer `call` signature, including how Keras decides when
    to pass in a `training` arg., and the semantics applied to
    the first positional arg in the call signature.

  Caveats:
    * TF2 will not prune unused variable updates (or unused outputs). You may
      need to adjust your forward pass code to avoid computations or variable
      updates that you don't intend to use. (E.g. by adding a flag to the
      `forward_pass` call signature and branching on it).
    * Avoid Nesting variable creation in tf.function inside of `forward_pass`
      While the layer may safely be used from inside a `tf.function`, using
      a function inside of `forward_pass` will break the variable scoping.
    * If you would like to nest Keras layers/models or other
      `VariableScopeLayer`s directly in `forward_pass`, you need to
      assign them as attributes of your layer so that Keras's standard
      object-oriented weights and loss tracking will kick in.
      See the intro to modules, layers, and models
      [guide](https://www.tensorflow.org/guide/intro_to_modules) for more info
  """

  @property
  @layer_utils.cached_per_instance
  def _call_full_argspec(self):
    # Argspec inspection is expensive and the call spec is used often, so it
    # makes sense to cache the result.
    return tf_inspect.getfullargspec(self.forward_pass)

  def forward_pass(self, *args, **kwargs):
    """Implement this method. It should include your model forward pass."""
    raise NotImplementedError

  @track_tf1_style_variables
  def call(self, *args, **kwargs):
    return self.forward_pass(*args, **kwargs)


@keras_export(v1=["keras.utils.get_or_create_layer"])
def get_or_create_layer(name, create_layer_method):
  """Use this method to track nested keras models in a shim-decorated method.

  This method can be used within a `tf.keras.Layer`'s methods decorated by
  the`track_tf1_style_variables` shim, to additionally track inner keras Model
  objects created within the same method. The inner model's variables and losses
  will be accessible via the outer model's `variables` and `losses` attributes.

  This enables tracking of inner keras models using TF2 behaviors, with minimal
  changes to existing TF1-style code.

  Example:

  ```python
  class NestedLayer(tf.keras.layers.Layer):

    def __init__(self, units, *args, **kwargs):
      super().__init__(*args, **kwargs)
      self.units = units

    def build_model(self):
      inp = tf.keras.Input(shape=(5, 5))
      dense_layer = tf.keras.layers.Dense(
          10, name="dense", kernel_regularizer="l2",
          kernel_initializer=tf.compat.v1.ones_initializer())
      model = tf.keras.Model(inputs=inp, outputs=dense_layer(inp))
      return model

    @tf.compat.v1.keras.utils.track_tf1_style_variables
    def call(self, inputs):
      model = tf.compat.v1.keras.utils.get_or_create_layer(
          "dense_model", self.build_model)
      return model(inputs)
  ```
  The inner model creation should be confined to its own zero-arg function,
  which should be passed into this method. In TF1, this method will immediately
  create and return the desired model, without any tracking.

  Args:
    name: A name to give the nested layer to track.
    create_layer_method: a Callable that takes no args and returns the nested
    layer.

  Returns:
    The created layer.
  """
  store = vs._get_default_variable_store()  # pylint: disable=protected-access
  if not isinstance(store, _EagerVariableStore):
    if not tf.compat.v1.executing_eagerly_outside_functions():
      # tf1 case; just create and return layer
      return create_layer_method()
    else:
      raise ValueError(
          "Tried to call get_or_create_layer in eager mode from a method not"
          "decorated with @tf.compat.v1.keras.utils.track_tf1_style_variables.")
  vs_name = tf.compat.v1.get_variable_scope().name
  name = f"{vs_name}/{name}"
  return store.get_or_create_layer(name, create_layer_method)
