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.

    Example:

        .. code-block:: python

            import tensorflow as tf
            enable_argscope_for_module(tf.layers)

    Args:
        log_shape (bool): print input/output shapes of each function.
    """
    if is_tfv2() and module == tf.layers:
        module = tf.compat.v1.layers
    for name, obj in getmembers(module):
        if isfunction(obj):
            setattr(module, name, enable_argscope_for_function(obj,
                    log_shape=log_shape))