<s> import numpy <EOL> from chainer import cuda <EOL> from chainer import function <EOL> from chainer . utils import type_check <EOL> class Sum ( function . Function ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , axis = None ) : <EOL> if axis is None : <EOL> self . axis = None <EOL> elif isinstance ( axis , int ) : <EOL> self . axis = ( axis , ) <EOL> elif isinstance ( axis , tuple ) and all ( isinstance ( a , int ) for a in axis ) : <EOL> if len ( set ( axis ) ) != len ( axis ) : <EOL> raise ValueError ( '<STR_LIT>' . format ( <EOL> '<STR_LIT:U+002CU+0020>' . join ( map ( str , axis ) ) ) ) <EOL> self . axis = axis <EOL> else : <EOL> raise TypeError ( '<STR_LIT>' ) <EOL> def check_type_forward ( self , in_types ) : <EOL> type_check . expect ( <EOL> in_types . size ( ) == <NUM_LIT:1> , <EOL> in_types [ <NUM_LIT:0> ] . dtype == numpy . float32 <EOL> ) <EOL> if self . axis is not None : <EOL> for axis in self . axis : <EOL> if axis >= <NUM_LIT:0> : <EOL> type_check . expect ( <EOL> axis < in_types [ <NUM_LIT:0> ] . ndim , <EOL> ) <EOL> else : <EOL> type_check . expect ( <EOL> - axis - <NUM_LIT:1> < in_types [ <NUM_LIT:0> ] . ndim , <EOL> ) <EOL> def forward ( self , x ) : <EOL> xp = cuda . get_array_module ( * x ) <EOL> return xp . asarray ( x [ <NUM_LIT:0> ] . sum ( axis = self . axis ) ) , <EOL> def backward ( self , x , gy ) : <EOL> xp = cuda . get_array_module ( * x ) <EOL> gx = xp . empty_like ( x [ <NUM_LIT:0> ] ) <EOL> if self . axis is None : <EOL> gx [ : ] = gy [ <NUM_LIT:0> ] <EOL> else : <EOL> gy = gy [ <NUM_LIT:0> ] <EOL> actual_axis = [ ] <EOL> for axis in self . axis : <EOL> if axis < <NUM_LIT:0> : <EOL> axis = len ( gx . shape ) + axis <EOL> actual_axis . append ( axis ) <EOL> for axis in sorted ( actual_axis ) : <EOL> gy = xp . expand_dims ( gy , axis = axis ) <EOL> gx [ : ] = gy <EOL> return gx , <EOL> def sum ( x , axis = None ) : <EOL> """<STR_LIT>""" <EOL> return Sum ( axis ) ( x ) </s>
