<commit_before># -*- coding: utf-8 -*-
# File: argscope.py

import copy
from collections import defaultdict
from contextlib import contextmanager
from functools import wraps
from inspect import getmembers, isfunction

from ..utils import logger
from .tower import get_current_tower_context

__all__ = ['argscope', 'get_arg_scope', 'enable_argscope_for_module',
           'enable_argscope_for_function']

_ArgScopeStack = []


@contextmanager
def argscope(layers, **kwargs):
    """
    Args:
        layers (list or layer): layer or list of layers to apply the arguments.

    Returns:
        a context where all appearance of these layer will by default have the
        arguments specified by kwargs.

    Example:
        .. code-block:: python

            with argscope(Conv2D, kernel_shape=3, nl=tf.nn.relu, out_channel=32):
                x = Conv2D('conv0', x)
                x = Conv2D('conv1', x)
                x = Conv2D('conv2', x, out_channel=64)  # override argscope

    """
    if not isinstance(layers, list):
        layers = [layers]

    # def _check_args_exist(l):
    #     args = inspect.getargspec(l).args
    #     for k, v in six.iteritems(kwargs):
    #         assert k in args, "No argument {} in {}".format(k, l.__name__)

    for l in layers:
        assert hasattr(l, 'symbolic_function'), "{} is not a registered layer".format(l.__name__)
        # _check_args_exist(l.symbolic_function)

    new_scope = copy.copy(get_arg_scope())
    for l in layers:
        new_scope[l.__name__].update(kwargs)
    _ArgScopeStack.append(new_scope)
    yield
    del _ArgScopeStack[-1]


def get_arg_scope():
    """
    Returns:
        dict: the current argscope.

    An argscope is a dict of dict: ``dict[layername] = {arg: val}``
    """
    if len(_ArgScopeStack) > 0:
        return _ArgScopeStack[-1]
    else:
        return defaultdict(dict)


def enable_argscope_for_function(func, log_shape=True):
    """Decorator for function to support argscope

    Example:

        .. code-block:: python

            from mylib import myfunc
            myfunc = enable_argscope_for_function(myfunc)

    Args:
        func: function which should be decorated.
        log_shape (bool): print input/output shapes of each function.

    Returns:
        The decorated function.

    """
    @wraps(func)
    def wrapped_func(*args, **kwargs):
        actual_args = copy.copy(get_arg_scope()[func.__name__])
        actual_args.update(kwargs)
        out_tensor = func(*args, **actual_args)
        in_tensor = args[0]

        ctx = get_current_tower_context()
        name = func.__name__ if 'name' not in kwargs else kwargs['name']
        if log_shape:
            if ('tower' not in ctx.ns_name.lower()) or ctx.is_main_training_tower:
                logger.info('%20s: %20s -> %20s' %
                            (name, in_tensor.shape.as_list(),
                             out_tensor.shape.as_list()))

        return out_tensor
    # argscope requires this property
    wrapped_func.symbolic_function = None
    return wrapped_func


def enable_argscope_for_module(module, log_shape=True):
    """
    Overwrite all functions of a given module to support argscope.
    Note that this function monkey-patches the module and therefore could
    have unexpected consequences.
    It has been only tested to work well with `tf.layers` module.
<commit_msg>argscope_for_function had some issues when a layer produces multiple … (#1049)<commit_after># -*- coding: utf-8 -*-
# File: argscope.py

import copy
from collections import defaultdict
from contextlib import contextmanager
from functools import wraps
from inspect import getmembers, isfunction

from ..utils import logger
from .tower import get_current_tower_context

__all__ = ['argscope', 'get_arg_scope', 'enable_argscope_for_module',
           'enable_argscope_for_function']

_ArgScopeStack = []


@contextmanager
def argscope(layers, **kwargs):
    """
    Args:
        layers (list or layer): layer or list of layers to apply the arguments.

    Returns:
        a context where all appearance of these layer will by default have the
        arguments specified by kwargs.

    Example:
        .. code-block:: python

            with argscope(Conv2D, kernel_shape=3, nl=tf.nn.relu, out_channel=32):
                x = Conv2D('conv0', x)
                x = Conv2D('conv1', x)
                x = Conv2D('conv2', x, out_channel=64)  # override argscope

    """
    if not isinstance(layers, list):
        layers = [layers]

    # def _check_args_exist(l):
    #     args = inspect.getargspec(l).args
    #     for k, v in six.iteritems(kwargs):
    #         assert k in args, "No argument {} in {}".format(k, l.__name__)

    for l in layers:
        assert hasattr(l, 'symbolic_function'), "{} is not a registered layer".format(l.__name__)
        # _check_args_exist(l.symbolic_function)

    new_scope = copy.copy(get_arg_scope())
    for l in layers:
        new_scope[l.__name__].update(kwargs)
    _ArgScopeStack.append(new_scope)
    yield
    del _ArgScopeStack[-1]


def get_arg_scope():
    """
    Returns:
        dict: the current argscope.

    An argscope is a dict of dict: ``dict[layername] = {arg: val}``
    """
    if len(_ArgScopeStack) > 0:
        return _ArgScopeStack[-1]
    else:
        return defaultdict(dict)


def enable_argscope_for_function(func, log_shape=True):
    """Decorator for function to support argscope

    Example:

        .. code-block:: python

            from mylib import myfunc
            myfunc = enable_argscope_for_function(myfunc)

    Args:
        func: A function mapping one or multiple tensors to one or multiple
            tensors.
        log_shape (bool): Specify whether the first input resp. output tensor
            shape should be printed once.

    Remarks:
        If the function `func` returns multiple input or output tensors,
        only the first input/output tensor shape is displayed during logging.

    Returns:
        The decorated function.

    """

    assert callable(func), "func should be a callable"

    @wraps(func)
    def wrapped_func(*args, **kwargs):
        actual_args = copy.copy(get_arg_scope()[func.__name__])
        actual_args.update(kwargs)
        out_tensor = func(*args, **actual_args)
        in_tensor = args[0]

        ctx = get_current_tower_context()
        name = func.__name__ if 'name' not in kwargs else kwargs['name']
        if log_shape:
            if ('tower' not in ctx.ns_name.lower()) or ctx.is_main_training_tower:
                # we assume the first parameter is the most interesting
                if isinstance(out_tensor, tuple):
                    out_tensor_descr = out_tensor[0]
                else:
                    out_tensor_descr = out_tensor
                logger.info('%20s: %20s -> %20s' %
                            (name, in_tensor.shape.as_list(),
                             out_tensor_descr.shape.as_list()))

        return out_tensor
    # argscope requires this property
    wrapped_func.symbolic_function = None
    return wrapped_func


def enable_argscope_for_module(module, log_shape=True):
    """
    Overwrite all functions of a given module to support argscope.
    Note that this function monkey-patches the module and therefore could
    have unexpected consequences.
    It has been only tested to work well with `tf.layers` module.

{'fraction_non_alphanumeric': 0.09164529626130381, 'fraction_numerical': 0.006883520043190714, 'mean_word_length': 2.5285714285714285, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '4867973', 'n_tokens_mistral': 2209, 'n_tokens_neox': 2142, 'n_words': 647}