<s> from __future__ import absolute_import , print_function <EOL> import sys <EOL> import warnings <EOL> import copy <EOL> import operator as op <EOL> import types <EOL> import math <EOL> import re <EOL> import numpy as np <EOL> from functools import partial <EOL> from future . utils import iteritems , raise_from , raise_ <EOL> from collections import defaultdict <EOL> def grad ( fun , argnum = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> @ attach_name_and_doc ( fun , argnum , '<STR_LIT>' ) <EOL> def gradfun ( * args , ** kwargs ) : <EOL> return backward_pass ( * forward_pass ( fun , args , kwargs , argnum ) ) <EOL> return gradfun <EOL> def forward_pass ( fun , args , kwargs , argnum = <NUM_LIT:0> ) : <EOL> tape = CalculationTape ( ) <EOL> arg_wrt = args [ argnum ] <EOL> start_node = new_node ( safe_type ( getval ( arg_wrt ) ) , [ tape ] ) <EOL> args = list ( args ) <EOL> args [ argnum ] = merge_tapes ( start_node , arg_wrt ) <EOL> try : end_node = fun ( * args , ** kwargs ) <EOL> except Exception as e : add_extra_error_message ( e ) <EOL> return start_node , end_node , tape <EOL> def backward_pass ( start_node , end_node , tape ) : <EOL> if not isinstance ( end_node , Node ) or tape not in end_node . tapes : <EOL> warnings . warn ( "<STR_LIT>" ) <EOL> return zeros_like ( start_node ) <EOL> if type ( end_node ) is not FloatNode : <EOL> try : <EOL> end_node = FloatNode . cast ( end_node , <NUM_LIT:1.0> ) <EOL> except TypeError : <EOL> raise TypeError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" . format ( type ( end_node . value ) ) ) <EOL> for node in tape : <EOL> node . outgrads = [ ] <EOL> end_node . tapes [ tape ] . outgrads = [ <NUM_LIT:1.0> ] <EOL> tape . complete = True <EOL> for node in tape [ : : - <NUM_LIT:1> ] : <EOL> if node . outgrads : <EOL> cur_outgrad = node . sum_outgrads ( ) <EOL> assert type ( new_node ( getval ( cur_outgrad ) ) ) == node . node_type , "<STR_LIT>" . format ( type ( new_node ( getval ( cur_outgrad ) ) ) , node . node_type ) <EOL> for gradfun , parent in node . parent_grad_ops : <EOL> og = cast_to_node_type ( gradfun ( cur_outgrad ) , parent . node_type , parent . node_value ) <EOL> parent . outgrads . append ( og ) <EOL> return cur_outgrad <EOL> def attach_name_and_doc ( fun , argnum , opname ) : <EOL> namestr = "<STR_LIT>" . format ( <EOL> op = opname . lower ( ) , fun = fun . __name__ , argnum = argnum ) <EOL> docstr = "<STR_LIT>" "<STR_LIT>" "<STR_LIT>" . format ( op = opname , fun = fun . __name__ , argnum = argnum ) <EOL> def wrap ( gradfun ) : <EOL> try : <EOL> gradfun . __name__ = namestr <EOL> gradfun . __doc__ = docstr <EOL> finally : <EOL> return gradfun <EOL> return wrap <EOL> def cast_to_node_type ( x , node_type , example ) : <EOL> if type ( new_node ( getval ( x ) ) ) is not node_type : <EOL> return node_type . cast ( x , example ) <EOL> else : <EOL> return x <EOL> class primitive ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , fun ) : <EOL> self . fun = fun <EOL> self . grads = { } <EOL> self . zero_grads = set ( ) <EOL> self . __name__ = fun . __name__ <EOL> self . __doc__ = fun . __doc__ <EOL> def gradmaker ( self , argnum , ans , args , kwargs ) : <EOL> try : <EOL> return self . grads [ argnum ] ( ans , * args , ** kwargs ) <EOL> except KeyError : <EOL> def error ( * args , ** kwargs ) : <EOL> if self . grads == { } : <EOL> errstr = "<STR_LIT>" <EOL> else : <EOL> errstr = "<STR_LIT>" <EOL> raise NotImplementedError ( errstr . format ( self . fun . __name__ , argnum ) ) <EOL> return error <EOL> def defgrad ( self , gradmaker , argnum = <NUM_LIT:0> ) : <EOL> self . grads [ argnum ] = gradmaker <EOL> def defgrads ( self , gradmaker , argnums ) : <EOL> for argnum in argnums : <EOL> self . defgrad ( partial ( gradmaker , argnum ) , argnum ) <EOL> def defgrad_is_zero ( self , argnums = ( <NUM_LIT:0> , ) ) : <EOL> for argnum in argnums : <EOL> self . zero_grads . add ( argnum ) <EOL> def __call__ ( self , * args , ** kwargs ) : <EOL> argvals = list ( args ) <EOL> ops = [ ] <EOL> tapes = set ( ) <EOL> for i , arg in enumerate ( args ) : <EOL> if isinstance ( arg , Node ) : <EOL> argvals [ i ] = arg . value <EOL> if i in self . zero_grads : continue <EOL> for tape , parent_rnode in iteritems ( arg . tapes ) : <EOL> if not tape . complete : <EOL> ops . append ( ( tape , i , parent_rnode ) ) <EOL> tapes . add ( tape ) <EOL> result = self . fun ( * argvals , ** kwargs ) <EOL> if result is NotImplemented : return result <EOL> if ops : <EOL> result = new_node ( result , tapes ) <EOL> for tape , argnum , parent in ops : <EOL> gradfun = self . gradmaker ( argnum , result , args , kwargs ) <EOL> rnode = result . tapes [ tape ] <EOL> rnode . parent_grad_ops . append ( ( gradfun , parent ) ) <EOL> return result <EOL> if sys . version_info >= ( <NUM_LIT:3> , ) : <EOL> def __get__ ( self , obj , objtype ) : <EOL> return types . MethodType ( self , obj ) <EOL> else : <EOL> def __get__ ( self , obj , objtype ) : <EOL> return types . MethodType ( self , obj , objtype ) <EOL> class primitive_with_aux ( primitive ) : <EOL> '''<STR_LIT>''' <EOL> def __call__ ( self , * args , ** kwargs ) : <EOL> argvals = list ( args ) <EOL> ops = [ ] <EOL> tapes = set ( ) <EOL> for i , arg in enumerate ( args ) : <EOL> if isinstance ( arg , Node ) : <EOL> argvals [ i ] = arg . value <EOL> if i in self . zero_grads : continue <EOL> for tape , parent_rnode in iteritems ( arg . tapes ) : <EOL> if not tape . complete : <EOL> ops . append ( ( tape , i , parent_rnode ) ) <EOL> tapes . add ( tape ) <EOL> result , aux = self . fun ( * argvals , ** kwargs ) <EOL> if result is NotImplemented : return result <EOL> if ops : <EOL> result = new_node ( result , tapes ) <EOL> for tape , argnum , parent in ops : <EOL> gradfun = self . gradmaker ( argnum , aux , result , args , kwargs ) <EOL> rnode = result . tapes [ tape ] <EOL> rnode . parent_grad_ops . append ( ( gradfun , parent ) ) <EOL> return result <EOL> def gradmaker ( self , argnum , aux , ans , args , kwargs ) : <EOL> try : <EOL> return self . grads [ argnum ] ( aux , ans , * args , ** kwargs ) <EOL> except KeyError : <EOL> return super ( primitive_with_aux , self ) . gradmaker ( argnum , ans , args , kwargs ) <EOL> @ primitive <EOL> def merge_tapes ( x , y ) : return x <EOL> merge_tapes . defgrad ( lambda ans , x , y : lambda g : g ) <EOL> merge_tapes . defgrad ( lambda ans , x , y : lambda g : g , argnum = <NUM_LIT:1> ) <EOL> def new_node ( value , tapes = [ ] ) : <EOL> try : <EOL> return Node . type_mappings [ type ( value ) ] ( value , tapes ) <EOL> except KeyError : <EOL> return NoDerivativeNode ( value , tapes ) <EOL> def zeros_like ( value ) : <EOL> if isinstance ( value , Node ) : <EOL> return value . zeros_like ( value ) <EOL> else : <EOL> return new_node ( value , [ ] ) . zeros_like ( value ) <EOL> class ReverseNode ( object ) : <EOL> __slots__ = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> def __init__ ( self , node_type , node_value ) : <EOL> self . parent_grad_ops = [ ] <EOL> self . outgrads = [ ] <EOL> self . node_type = node_type <EOL> self . node_value = node_value <EOL> def sum_outgrads ( self ) : <EOL> return self . node_type . sum_outgrads ( self . outgrads ) <EOL> class Node ( object ) : <EOL> __slots__ = [ '<STR_LIT:value>' , '<STR_LIT>' ] <EOL> Rnode = ReverseNode <EOL> type_mappings = { } <EOL> def __init__ ( self , value , tapes ) : <EOL> self . value = value <EOL> self . tapes = { } <EOL> for tape in tapes : <EOL> new_rnode = self . Rnode ( type ( self ) , value ) <EOL> tape . append ( new_rnode ) <EOL> self . tapes [ tape ] = new_rnode <EOL> def __bool__ ( self ) : <EOL> return bool ( self . value ) <EOL> __nonzero__ = __bool__ <EOL> @ staticmethod <EOL> def sum_outgrads ( outgrads ) : <EOL> return sum ( outgrads [ <NUM_LIT:1> : ] , outgrads [ <NUM_LIT:0> ] ) <EOL> def __str__ ( self ) : <EOL> return "<STR_LIT>" . format ( <EOL> type ( self ) . __name__ , str ( self . value ) , len ( self . tapes ) ) <EOL> @ primitive <EOL> def cast ( value , caster ) : <EOL> return caster ( value ) <EOL> cast . defgrad ( lambda * args : I ) <EOL> getval = lambda x : x . value if isinstance ( x , Node ) else x <EOL> class CalculationTape ( list ) : <EOL> def __init__ ( self ) : <EOL> self . complete = False <EOL> def __hash__ ( self ) : <EOL> return id ( self ) <EOL> class FloatNode ( Node ) : <EOL> __slots__ = [ ] <EOL> @ staticmethod <EOL> def zeros_like ( value ) : <EOL> return <NUM_LIT:0.0> <EOL> @ staticmethod <EOL> def cast ( value , example ) : <EOL> return cast ( value , cast_to_float ) <EOL> Node . type_mappings [ float ] = FloatNode <EOL> def cast_to_float ( x ) : <EOL> if np . iscomplexobj ( x ) : <EOL> x = np . real ( x ) <EOL> return float ( x ) <EOL> class ComplexNode ( FloatNode ) : <EOL> @ staticmethod <EOL> def zeros_like ( value ) : <EOL> return <NUM_LIT:0.0> + <NUM_LIT> <EOL> @ staticmethod <EOL> def cast ( value , example ) : <EOL> return cast ( value , cast_to_complex ) <EOL> def cast_to_complex ( value ) : <EOL> if isinstance ( value , np . ndarray ) : <EOL> return complex ( value [ ( ) ] ) <EOL> else : <EOL> return complex ( value ) <EOL> Node . type_mappings [ complex ] = ComplexNode <EOL> def safe_type ( value ) : <EOL> if isinstance ( value , int ) : <EOL> warnings . warn ( "<STR_LIT>" ) <EOL> return float ( value ) <EOL> else : <EOL> return value <EOL> if sys . version_info >= ( <NUM_LIT:3> , ) : <EOL> DIV = '<STR_LIT>' <EOL> RDIV = '<STR_LIT>' <EOL> else : <EOL> DIV = '<STR_LIT>' <EOL> RDIV = '<STR_LIT>' <EOL> differentiable_ops = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , DIV , RDIV ] <EOL> nondifferentiable_ops = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , ] <EOL> for float_op in differentiable_ops + nondifferentiable_ops : <EOL> setattr ( FloatNode , float_op , primitive ( getattr ( float , float_op ) ) ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . defgrad ( lambda ans , x : op . neg ) <EOL> for comp_op in nondifferentiable_ops : <EOL> FloatNode . __dict__ [ comp_op ] . defgrad_is_zero ( argnums = ( <NUM_LIT:0> , <NUM_LIT:1> ) ) <EOL> I = lambda g : g <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . defgrad ( lambda ans , x , y : I ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . defgrad ( lambda ans , x , y : I , argnum = <NUM_LIT:1> ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . defgrad ( lambda ans , x , y : lambda g : y * g ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . defgrad ( lambda ans , x , y : lambda g : x * g , argnum = <NUM_LIT:1> ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . defgrad ( lambda ans , x , y : I ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . defgrad ( lambda ans , x , y : op . neg , argnum = <NUM_LIT:1> ) <EOL> FloatNode . __dict__ [ DIV ] . defgrad ( lambda ans , x , y : lambda g : g / y ) <EOL> FloatNode . __dict__ [ DIV ] . defgrad ( lambda ans , x , y : lambda g : - g * x / y ** <NUM_LIT:2> , argnum = <NUM_LIT:1> ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . defgrad ( lambda ans , x , y : lambda g : g * y * x ** ( y - <NUM_LIT:1> ) ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . defgrad ( lambda ans , x , y : lambda g : g * log ( x ) * x ** y , argnum = <NUM_LIT:1> ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . defgrad ( lambda ans , x , y : I ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . defgrad ( lambda ans , x , y : lambda g : - g * floor ( x / y ) , argnum = <NUM_LIT:1> ) <EOL> log = primitive ( math . log ) <EOL> log . defgrad ( lambda ans , x : lambda g : g / x ) <EOL> floor = primitive ( math . floor ) <EOL> floor . defgrad_is_zero ( ) <EOL> def swap_args ( grads ) : <EOL> grad_0 , grad_1 = grads [ <NUM_LIT:1> ] , grads [ <NUM_LIT:0> ] <EOL> return { <NUM_LIT:0> : lambda ans , y , x : grad_0 ( ans , x , y ) , <EOL> <NUM_LIT:1> : lambda ans , y , x : grad_1 ( ans , x , y ) } <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . grads = swap_args ( FloatNode . __dict__ [ '<STR_LIT>' ] . grads ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . grads = swap_args ( FloatNode . __dict__ [ '<STR_LIT>' ] . grads ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . grads = swap_args ( FloatNode . __dict__ [ '<STR_LIT>' ] . grads ) <EOL> FloatNode . __dict__ [ RDIV ] . grads = swap_args ( FloatNode . __dict__ [ DIV ] . grads ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . grads = swap_args ( FloatNode . __dict__ [ '<STR_LIT>' ] . grads ) <EOL> FloatNode . __dict__ [ '<STR_LIT>' ] . grads = swap_args ( FloatNode . __dict__ [ '<STR_LIT>' ] . grads ) <EOL> class NoDerivativeReverseNode ( ReverseNode ) : <EOL> def __init__ ( self , node_type , node_value ) : <EOL> super ( NoDerivativeReverseNode , self ) . __init__ ( node_type , node_value ) <EOL> self . type = type ( node_value ) <EOL> def sum_outgrads ( self ) : <EOL> raise TypeError ( "<STR_LIT>" . format ( self . type ) ) <EOL> class NoDerivativeNode ( FloatNode ) : <EOL> Rnode = NoDerivativeReverseNode <EOL> @ staticmethod <EOL> def cast ( value , example ) : <EOL> return example <EOL> class AutogradHint ( Exception ) : <EOL> def __init__ ( self , message , subexception_type = None , subexception_val = None ) : <EOL> self . message = message <EOL> self . subexception_type = subexception_type <EOL> self . subexception_val = subexception_val <EOL> def __str__ ( self ) : <EOL> if self . subexception_type : <EOL> return '<STR_LIT>' . format ( <EOL> message = self . message , <EOL> name = self . subexception_type . __name__ , <EOL> str = self . subexception_type ( self . subexception_val ) ) <EOL> else : <EOL> return self . message <EOL> common_errors = [ <EOL> ( ( TypeError , r'<STR_LIT>' ) , <EOL> "<STR_LIT>" ) , <EOL> ( ( TypeError , r"<STR_LIT>" ) , <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) , <EOL> ] <EOL> def check_common_errors ( error_type , error_message ) : <EOL> keys , vals = zip ( * common_errors ) <EOL> match = lambda key : error_type == key [ <NUM_LIT:0> ] and len ( re . findall ( key [ <NUM_LIT:1> ] , error_message ) ) != <NUM_LIT:0> <EOL> matches = map ( match , keys ) <EOL> num_matches = sum ( matches ) <EOL> if num_matches == <NUM_LIT:1> : <EOL> return vals [ matches . index ( True ) ] <EOL> def add_extra_error_message ( e ) : <EOL> etype , value , traceback = sys . exc_info ( ) <EOL> extra_message = check_common_errors ( type ( e ) , str ( e ) ) <EOL> if extra_message : <EOL> if sys . version_info >= ( <NUM_LIT:3> , ) : <EOL> raise_from ( AutogradHint ( extra_message ) , e ) <EOL> else : <EOL> raise_ ( AutogradHint , ( extra_message , etype , value ) , traceback ) <EOL> raise_ ( etype , value , traceback ) </s>
