tools/autograd/load_derivatives.py
def process_definition(defn, declarations_by_signature):
    """Processes a single entry `defn` in derivatives.yaml"""

    def canonical_declaration(declarations, name):
        for declaration in declarations:
            if declaration['name'] == name:
                return declaration
        # some functions only have in-place variants
        assert name + '_' == declarations[0]['name']
        return declarations[0]

    def split_names(raw_names):
        """Given "foo, bar", return ["foo", "bar"]."""
        return [x.strip() for x in raw_names.split(',')]

    def lookup_pred(pred, xs):
        """Return the index of the first element of xs matching pred."""
        return next((i, x) for i, x in enumerate(xs) if pred(x))

    def check_grad_usage(defn_name, declaration, derivatives):
        """
        Check for some subtle mistakes one might make when writing gradients.
        These mistakes will compile, but will be latent until a function is
        used with double backwards.
        """

        used_grad = 0
        used_grads = 0
        fully_implemented = True
        used_grads_indices = []
        for d in derivatives:
            formula = d['formula']
            used_grad += len(re.findall(IDENT_REGEX.format('grad'), formula))
            used_grads += len(re.findall(IDENT_REGEX.format('grads'), formula))
            fully_implemented = \
                fully_implemented and \
                not re.search(IDENT_REGEX.format('not_implemented'), formula)
            used_grads_indices.extend(used_gradient_indices(formula))
        assert used_grads >= len(used_grads_indices)
        only_used_grads_indices = used_grads == len(used_grads_indices)

        if used_grad and used_grads:
            raise RuntimeError("Derivative definition of {} in derivatives.yaml illegally "
                               "mixes use of 'grad' and 'grads'. Consider replacing "
                               "occurrences of 'grad' with 'grads[0]'".format(defn_name))

        if only_used_grads_indices and set(used_grads_indices) == {0}:
            raise RuntimeError("Derivative definition of {} in derivatives.yaml solely "
                               "refers to 'grads[0]'.  If the first output is indeed the "
                               "only differentiable output, replace 'grads[0]' with 'grad'; "
                               "otherwise, there is a likely error in your derivatives "
                               "declaration.".format(defn_name))

        hardcoded_diff = HARDCODED_DIFFERENTIABLE_OUTPUTS.get(defn_name)
        if hardcoded_diff:
            if used_grad:
                raise RuntimeError("Derivative definition {} has hard-coded differentiable "
                                   "outputs in gen_autograd.py, but used grad (which implies "
                                   "only the first output is differentiable) in its "
                                   "derivative declaration.  You likely meant to write "
                                   "grads[i] for some i instead.".format(defn_name))
            if only_used_grads_indices and set(used_grads_indices) != set(hardcoded_diff):
                raise RuntimeError("Derivative definition {} has hard-coded differentiable "
                                   "outputs {}, but the used grads in the derivative "
                                   "definitions are only {}.  Either your derivatives "
                                   "declaration is wrong, or the value of "
                                   "HARDCODED_DIFFERENTIABLE_OUTPUTS in gen_autograd.py "
                                   "is wrong.".format(defn_name, hardcoded_diff,
                                                      used_grads_indices))
        else:
            if fully_implemented and not used_grad and \
               used_grads and only_used_grads_indices and \
               set(used_grads_indices) != set(range(len(declaration['returns']))):
                raise RuntimeError("Derivative definition of {} in derivatives.yaml does "
                                   "not refer to the gradients of all of its outputs.  Either "
                                   "the derivatives declaration is wrong, OR you have some "
                                   "non-differentiable outputs.  If you have a single "
                                   "differentiable output, make it the first output in ATen "
                                   "and reference its gradient with 'grad'; otherwise, hard "
                                   "code the list of differentiable outputs in "
                                   "HARDCODED_DIFFERENTIABLE_OUTPUTS in gen_autograd.py."
                                   .format(defn_name))

    def set_up_derivatives(defn_name, defn, declaration):
        # Determine the set of inputs which have gradients
        args_with_gradients_set = set()
        for raw_names in defn:
            args_with_gradients_set |= set(split_names(raw_names))

        # Next, let us determine the list of inputs in order.
        args_with_gradients = []
        for arg in declaration['arguments']:
            if arg['name'] not in args_with_gradients_set:
                continue
            args_with_gradients.append(arg)

        # Set up the derivative information
        derivatives = []
        for raw_names in sorted(defn.keys()):
            formula = defn[raw_names]
            names = split_names(raw_names)
            derivatives.append(create_derivative(declaration, formula, names))

        # Test to see if the use of 'grads' makes sense.
        check_grad_usage(defn_name, declaration, derivatives)

        return derivatives, args_with_gradients

    def unzip(xs):
        return zip(*xs)

    # NB: Removes 'name' from defn dictionary
    defn_name, params = split_name_params(defn.pop('name'))
    param_types, param_names = unzip([p.split(' ') for p in params if p != '*'])
    if 'grad_input_mask' in param_names:
        raise RuntimeError("Signature for {} has an argument named grad_input_mask, "
                           "but this name would be shadowed by our codegen. "
                           "Please use a different name in Declarations.cwrap."
                           .format(defn_name))
    signature = '{}({})'.format(defn_name, ', '.join(param_types))

    declarations = declarations_by_signature[signature]
    if len(declarations) == 0:
        avail = [k for k, v in declarations_by_signature.items()
                 if k.startswith(defn_name + '(') and len(v) > 0]
        raise RuntimeError('no ATen declaration found for: {}.  '
                           'Available signatures: {}'.format(signature, ', '.join(avail)))
    canonical = canonical_declaration(declarations, defn_name)

    # TODO: Check the types line up
    if len(param_names) != len(canonical['args']):
        raise RuntimeError('Signature for {} has {} arguments ({}), but '
                           'Declarations.yaml records {} arguments ({})'
                           .format(defn_name,
                                   len(param_names),
                                   ', '.join(param_names),
                                   len(canonical['args']),
                                   ', '.join(canonical['args'])))
    for i, (x, y) in enumerate(zip(param_names, canonical['args'])):
        if x != y:
            raise RuntimeError('Argument {} of {} has different names in '
                               'derivatives.yaml ({}) and '
                               'Declarations.yaml ({})'
                               .format(i, defn_name, x, y))

    derivatives, args_with_gradients = set_up_derivatives(defn_name, defn, canonical)
    return create_autograd_function(defn_name, derivatives, args_with_gradients, signature, canonical)
Processes a single entry `defn` in derivatives.yaml
Processes a single entry `defn` in derivatives.yaml
Processes a single entry ` defn ` in derivatives . yaml
def process_definition(defn, declarations_by_signature):
    def canonical_declaration(declarations, name):
        for declaration in declarations:
            if declaration['name'] == name:
                return declaration
        assert name + '_' == declarations[0]['name']
        return declarations[0]
    def split_names(raw_names):
        return [x.strip() for x in raw_names.split(',')]
    def lookup_pred(pred, xs):
        return next((i, x) for i, x in enumerate(xs) if pred(x))
    def check_grad_usage(defn_name, declaration, derivatives):
        used_grad = 0
        used_grads = 0
        fully_implemented = True
        used_grads_indices = []
        for d in derivatives:
            formula = d['formula']
            used_grad += len(re.findall(IDENT_REGEX.format('grad'), formula))
            used_grads += len(re.findall(IDENT_REGEX.format('grads'), formula))
            fully_implemented = \
                fully_implemented and \
                not re.search(IDENT_REGEX.format('not_implemented'), formula)
            used_grads_indices.extend(used_gradient_indices(formula))
        assert used_grads >= len(used_grads_indices)
        only_used_grads_indices = used_grads == len(used_grads_indices)
        if used_grad and used_grads:
            raise RuntimeError("Derivative definition of {} in derivatives.yaml illegally "
                               "mixes use of 'grad' and 'grads'. Consider replacing "
                               "occurrences of 'grad' with 'grads[0]'".format(defn_name))
        if only_used_grads_indices and set(used_grads_indices) == {0}:
            raise RuntimeError("Derivative definition of {} in derivatives.yaml solely "
                               "refers to 'grads[0]'.  If the first output is indeed the "
                               "only differentiable output, replace 'grads[0]' with 'grad'; "
                               "otherwise, there is a likely error in your derivatives "
                               "declaration.".format(defn_name))
        hardcoded_diff = HARDCODED_DIFFERENTIABLE_OUTPUTS.get(defn_name)
        if hardcoded_diff:
            if used_grad:
                raise RuntimeError("Derivative definition {} has hard-coded differentiable "
                                   "outputs in gen_autograd.py, but used grad (which implies "
                                   "only the first output is differentiable) in its "
                                   "derivative declaration.  You likely meant to write "
                                   "grads[i] for some i instead.".format(defn_name))
            if only_used_grads_indices and set(used_grads_indices) != set(hardcoded_diff):
                raise RuntimeError("Derivative definition {} has hard-coded differentiable "
                                   "outputs {}, but the used grads in the derivative "
                                   "definitions are only {}.  Either your derivatives "
                                   "declaration is wrong, or the value of "
                                   "HARDCODED_DIFFERENTIABLE_OUTPUTS in gen_autograd.py "
                                   "is wrong.".format(defn_name, hardcoded_diff,
                                                      used_grads_indices))
        else:
            if fully_implemented and not used_grad and \
               used_grads and only_used_grads_indices and \
               set(used_grads_indices) != set(range(len(declaration['returns']))):
                raise RuntimeError("Derivative definition of {} in derivatives.yaml does "
                                   "not refer to the gradients of all of its outputs.  Either "
                                   "the derivatives declaration is wrong, OR you have some "
                                   "non-differentiable outputs.  If you have a single "
                                   "differentiable output, make it the first output in ATen "
                                   "and reference its gradient with 'grad'; otherwise, hard "
                                   "code the list of differentiable outputs in "
                                   "HARDCODED_DIFFERENTIABLE_OUTPUTS in gen_autograd.py."
                                   .format(defn_name))
    def set_up_derivatives(defn_name, defn, declaration):
        args_with_gradients_set = set()
        for raw_names in defn:
            args_with_gradients_set |= set(split_names(raw_names))
        args_with_gradients = []
        for arg in declaration['arguments']:
            if arg['name'] not in args_with_gradients_set:
                continue
            args_with_gradients.append(arg)
        derivatives = []
        for raw_names in sorted(defn.keys()):
            formula = defn[raw_names]
            names = split_names(raw_names)
            derivatives.append(create_derivative(declaration, formula, names))
        check_grad_usage(defn_name, declaration, derivatives)
        return derivatives, args_with_gradients
    def unzip(xs):
        return zip(*xs)
    defn_name, params = split_name_params(defn.pop('name'))
    param_types, param_names = unzip([p.split(' ') for p in params if p != '*'])
    if 'grad_input_mask' in param_names:
        raise RuntimeError("Signature for {} has an argument named grad_input_mask, "
                           "but this name would be shadowed by our codegen. "
                           "Please use a different name in Declarations.cwrap."
                           .format(defn_name))
    signature = '{}({})'.format(defn_name, ', '.join(param_types))
    declarations = declarations_by_signature[signature]
    if len(declarations) == 0:
        avail = [k for k, v in declarations_by_signature.items()
                 if k.startswith(defn_name + '(') and len(v) > 0]
        raise RuntimeError('no ATen declaration found for: {}.  '
                           'Available signatures: {}'.format(signature, ', '.join(avail)))
    canonical = canonical_declaration(declarations, defn_name)
    if len(param_names) != len(canonical['args']):
        raise RuntimeError('Signature for {} has {} arguments ({}), but '
                           'Declarations.yaml records {} arguments ({})'
                           .format(defn_name,
                                   len(param_names),
                                   ', '.join(param_names),
                                   len(canonical['args']),
                                   ', '.join(canonical['args'])))
    for i, (x, y) in enumerate(zip(param_names, canonical['args'])):
        if x != y:
            raise RuntimeError('Argument {} of {} has different names in '
                               'derivatives.yaml ({}) and '
                               'Declarations.yaml ({})'
                               .format(i, defn_name, x, y))
    derivatives, args_with_gradients = set_up_derivatives(defn_name, defn, canonical)
    return create_autograd_function(defn_name, derivatives, args_with_gradients, signature, canonical)
def process_definition ( defn , declarations_by_signature ) : def canonical_declaration ( declarations , name ) : for declaration in declarations : if declaration [ 'name' ] == name : return declaration assert name + '_' == declarations [ 0 ] [ 'name' ] return declarations [ 0 ] def split_names ( raw_names ) : """Given "foo, bar", return ["foo", "bar"].""" return [ x . strip ( ) for x in raw_names . split ( ',' ) ] def lookup_pred ( pred , xs ) : """Return the index of the first element of xs matching pred.""" return next ( ( i , x ) for i , x in enumerate ( xs ) if pred ( x ) ) def check_grad_usage ( defn_name , declaration , derivatives ) : """
        Check for some subtle mistakes one might make when writing gradients.
        These mistakes will compile, but will be latent until a function is
        used with double backwards.
        """ used_grad = 0 used_grads = 0 fully_implemented = True used_grads_indices = [ ] for d in derivatives : formula = d [ 'formula' ] used_grad += len ( re . findall ( IDENT_REGEX . format ( 'grad' ) , formula ) ) used_grads += len ( re . findall ( IDENT_REGEX . format ( 'grads' ) , formula ) ) fully_implemented = fully_implemented and not re . search ( IDENT_REGEX . format ( 'not_implemented' ) , formula ) used_grads_indices . extend ( used_gradient_indices ( formula ) ) assert used_grads >= len ( used_grads_indices ) only_used_grads_indices = used_grads == len ( used_grads_indices ) if used_grad and used_grads : raise RuntimeError ( "Derivative definition of {} in derivatives.yaml illegally " "mixes use of 'grad' and 'grads'. Consider replacing " "occurrences of 'grad' with 'grads[0]'" . format ( defn_name ) ) if only_used_grads_indices and set ( used_grads_indices ) == { 0 } : raise RuntimeError ( "Derivative definition of {} in derivatives.yaml solely " "refers to 'grads[0]'.  If the first output is indeed the " "only differentiable output, replace 'grads[0]' with 'grad'; " "otherwise, there is a likely error in your derivatives " "declaration." . format ( defn_name ) ) hardcoded_diff = HARDCODED_DIFFERENTIABLE_OUTPUTS . get ( defn_name ) if hardcoded_diff : if used_grad : raise RuntimeError ( "Derivative definition {} has hard-coded differentiable " "outputs in gen_autograd.py, but used grad (which implies " "only the first output is differentiable) in its " "derivative declaration.  You likely meant to write " "grads[i] for some i instead." . format ( defn_name ) ) if only_used_grads_indices and set ( used_grads_indices ) != set ( hardcoded_diff ) : raise RuntimeError ( "Derivative definition {} has hard-coded differentiable " "outputs {}, but the used grads in the derivative " "definitions are only {}.  Either your derivatives " "declaration is wrong, or the value of " "HARDCODED_DIFFERENTIABLE_OUTPUTS in gen_autograd.py " "is wrong." . format ( defn_name , hardcoded_diff , used_grads_indices ) ) else : if fully_implemented and not used_grad and used_grads and only_used_grads_indices and set ( used_grads_indices ) != set ( range ( len ( declaration [ 'returns' ] ) ) ) : raise RuntimeError ( "Derivative definition of {} in derivatives.yaml does " "not refer to the gradients of all of its outputs.  Either " "the derivatives declaration is wrong, OR you have some " "non-differentiable outputs.  If you have a single " "differentiable output, make it the first output in ATen " "and reference its gradient with 'grad'; otherwise, hard " "code the list of differentiable outputs in " "HARDCODED_DIFFERENTIABLE_OUTPUTS in gen_autograd.py." . format ( defn_name ) ) def set_up_derivatives ( defn_name , defn , declaration ) : args_with_gradients_set = set ( ) for raw_names in defn : args_with_gradients_set |= set ( split_names ( raw_names ) ) args_with_gradients = [ ] for arg in declaration [ 'arguments' ] : if arg [ 'name' ] not in args_with_gradients_set : continue args_with_gradients . append ( arg ) derivatives = [ ] for raw_names in sorted ( defn . keys ( ) ) : formula = defn [ raw_names ] names = split_names ( raw_names ) derivatives . append ( create_derivative ( declaration , formula , names ) ) check_grad_usage ( defn_name , declaration , derivatives ) return derivatives , args_with_gradients def unzip ( xs ) : return zip ( * xs ) defn_name , params = split_name_params ( defn . pop ( 'name' ) ) param_types , param_names = unzip ( [ p . split ( ' ' ) for p in params if p != '*' ] ) if 'grad_input_mask' in param_names : raise RuntimeError ( "Signature for {} has an argument named grad_input_mask, " "but this name would be shadowed by our codegen. " "Please use a different name in Declarations.cwrap." . format ( defn_name ) ) signature = '{}({})' . format ( defn_name , ', ' . join ( param_types ) ) declarations = declarations_by_signature [ signature ] if len ( declarations ) == 0 : avail = [ k for k , v in declarations_by_signature . items ( ) if k . startswith ( defn_name + '(' ) and len ( v ) > 0 ] raise RuntimeError ( 'no ATen declaration found for: {}.  ' 'Available signatures: {}' . format ( signature , ', ' . join ( avail ) ) ) canonical = canonical_declaration ( declarations , defn_name ) if len ( param_names ) != len ( canonical [ 'args' ] ) : raise RuntimeError ( 'Signature for {} has {} arguments ({}), but ' 'Declarations.yaml records {} arguments ({})' . format ( defn_name , len ( param_names ) , ', ' . join ( param_names ) , len ( canonical [ 'args' ] ) , ', ' . join ( canonical [ 'args' ] ) ) ) for i , ( x , y ) in enumerate ( zip ( param_names , canonical [ 'args' ] ) ) : if x != y : raise RuntimeError ( 'Argument {} of {} has different names in ' 'derivatives.yaml ({}) and ' 'Declarations.yaml ({})' . format ( i , defn_name , x , y ) ) derivatives , args_with_gradients = set_up_derivatives ( defn_name , defn , canonical ) return create_autograd_function ( defn_name , derivatives , args_with_gradients , signature , canonical )
Processes a single entry `defn` in derivatives.yaml
Processes a single entry ` defn ` in derivatives . yaml
"""Processes a single entry `defn` in derivatives.yaml""" # some functions only have in-place variants """Given "foo, bar", return ["foo", "bar"].""" """Return the index of the first element of xs matching pred.""" """
        Check for some subtle mistakes one might make when writing gradients.
        These mistakes will compile, but will be latent until a function is
        used with double backwards.
        """ # Determine the set of inputs which have gradients # Next, let us determine the list of inputs in order. # Set up the derivative information # Test to see if the use of 'grads' makes sense. # NB: Removes 'name' from defn dictionary # TODO: Check the types line up
defn
declarations_by_signature