def __init__(self, func, regex='.*'):
        """
        Args:
            func: a user-supplied function which takes one or two arguments.
                The argument(s) can be either a `grad` tensor, or `grad` and `var`.
                The function should return the new gradient to be used.
                If it 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.
        """
        args = inspect.getfullargspec(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__()

        Args:
            func: a user-supplied function which takes one or two arguments.
                The argument(s) can be either a `grad` tensor, or `grad` and `var`.
                The function should return the new gradient to be used.
                If it 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.
        
tensorpack/tfutils/gradproc.py