'Args:
func: takes a grad or (grad, var) pair and returns a grad. If return None, the
gradient is discarded (hence no update to the variable will happen).
regex (str): used to match variables. Defaults to match all variables.'

def __init__(self, func, regex='.*'):

    args = inspect.getargspec(func).args
    arg_num = (len(args) - inspect.ismethod(func))
    assert (arg_num in [1, 2]), 'The    function    must    take    1    or    2    arguments!        ({})'.format(args)
    if (arg_num == 1):
        self.func = (lambda grad, var: func(grad))
    else:
        self.func = func
    if (not regex.endswith('$')):
        regex = (regex + '$')
    self.regex = regex
    super(MapGradient, self).__init__()
