<s> import numpy as np <EOL> import scipy . linalg as spla <EOL> from . import Differentiable <EOL> class MatMult ( Differentiable ) : <EOL> __slots__ = [ '<STR_LIT:A>' , '<STR_LIT:B>' ] <EOL> def __init__ ( self , A , B , * args ) : <EOL> if len ( args ) > <NUM_LIT:0> : <EOL> B = MatMult ( B , * args ) <EOL> super ( MatMult , self ) . __init__ ( ( A , B ) ) <EOL> self . A = A <EOL> self . B = B <EOL> def _compute_value ( self ) : <EOL> A_val , B_val = self . A . value , self . B . value <EOL> if A_val . ndim > <NUM_LIT:2> or B_val . ndim > <NUM_LIT:2> : <EOL> raise Exception ( "<STR_LIT>" % ( self . A . shape ) ) <EOL> if A_val . shape [ - <NUM_LIT:1> ] != B_val . shape [ <NUM_LIT:0> ] : <EOL> raise Exception ( "<STR_LIT>" % ( self . A . shape , self . B . shape ) ) <EOL> return np . dot ( self . A . value , self . B . value ) <EOL> def _local_grad ( self , parent , d_out_d_self ) : <EOL> if parent == <NUM_LIT:0> : <EOL> B_val = self . B . value <EOL> if B_val . ndim == <NUM_LIT:2> : <EOL> return np . dot ( d_out_d_self , B_val . T ) <EOL> else : <EOL> return np . outer ( d_out_d_self , B_val ) <EOL> elif parent == <NUM_LIT:1> : <EOL> A_val = self . A . value <EOL> if A_val . ndim == <NUM_LIT:2> : <EOL> return np . dot ( A_val . T , d_out_d_self ) <EOL> else : <EOL> return np . outer ( A_val , d_out_d_self ) <EOL> else : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> class MatSum ( Differentiable ) : <EOL> __slots__ = [ '<STR_LIT:A>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> def __init__ ( self , A , axis = None , keepdims = True ) : <EOL> super ( MatSum , self ) . __init__ ( ( A , ) ) <EOL> if axis is not None and type ( axis ) != int : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> self . A = A <EOL> self . axis = axis <EOL> self . keepdims = keepdims <EOL> def _compute_value ( self ) : <EOL> return np . sum ( self . A . value , axis = self . axis , keepdims = self . keepdims ) <EOL> def _local_grad ( self , parent , d_out_d_self ) : <EOL> if not self . keepdims and self . axis is not None : <EOL> expanded_d_out_d_self = np . expand_dims ( d_out_d_self , self . axis ) <EOL> return expanded_d_out_d_self * np . ones ( self . A . shape ) <EOL> else : <EOL> return d_out_d_self * np . ones ( self . A . shape ) <EOL> class MatMean ( Differentiable ) : <EOL> __slots__ = [ '<STR_LIT:A>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> def __init__ ( self , A , axis = None , keepdims = True ) : <EOL> super ( MatMean , self ) . __init__ ( ( A , ) ) <EOL> if axis is not None and type ( axis ) != int : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> self . A = A <EOL> self . axis = axis <EOL> self . keepdims = keepdims <EOL> def _compute_value ( self ) : <EOL> return np . mean ( self . A . value , axis = self . axis , keepdims = self . keepdims ) <EOL> def _local_grad ( self , parent , d_out_d_self ) : <EOL> N = float ( self . A . value . size ) if self . axis is None else float ( self . A . shape [ self . axis ] ) <EOL> if not self . keepdims and self . axis is not None : <EOL> expanded_d_out_d_self = np . expand_dims ( d_out_d_self , self . axis ) <EOL> return expanded_d_out_d_self * <NUM_LIT:1.0> / N * np . ones ( self . A . shape ) <EOL> else : <EOL> return d_out_d_self * <NUM_LIT:1.0> / N * np . ones ( self . A . shape ) <EOL> class MatAdd ( Differentiable ) : <EOL> __slots__ = [ ] <EOL> def __init__ ( self , * args ) : <EOL> super ( MatAdd , self ) . __init__ ( args ) <EOL> def _compute_value ( self ) : <EOL> return sum ( [ p . value for p in self . _parents ] ) <EOL> def _local_grad ( self , parent , d_out_d_self ) : <EOL> parent_shape = self . _parents [ parent ] . shape <EOL> num_singletons = len ( d_out_d_self . shape ) - len ( parent_shape ) <EOL> if num_singletons > <NUM_LIT:0> : <EOL> extra_singletons = tuple ( range ( num_singletons ) ) <EOL> result = np . sum ( d_out_d_self , axis = extra_singletons , keepdims = False ) <EOL> else : <EOL> result = d_out_d_self <EOL> assert len ( result . shape ) == len ( parent_shape ) <EOL> original_singletons = tuple ( np . where ( np . array ( parent_shape ) == <NUM_LIT:1> ) [ <NUM_LIT:0> ] ) <EOL> return np . sum ( result , axis = original_singletons , keepdims = True ) <EOL> class MatElemMult ( Differentiable ) : <EOL> """<STR_LIT>""" <EOL> __slots__ = [ '<STR_LIT:A>' , '<STR_LIT:B>' ] <EOL> def __init__ ( self , A , B , * args ) : <EOL> if len ( args ) > <NUM_LIT:0> : <EOL> B = MatElemMult ( B , * args ) <EOL> super ( MatElemMult , self ) . __init__ ( ( A , B ) ) <EOL> self . A = A <EOL> self . B = B <EOL> def _compute_value ( self ) : <EOL> return self . A . value * self . B . value <EOL> def _local_grad ( self , parent , d_out_d_self ) : <EOL> """<STR_LIT>""" <EOL> parent_shape = self . _parents [ parent ] . shape <EOL> other_parent = <NUM_LIT:1> if parent == <NUM_LIT:0> else <NUM_LIT:0> <EOL> other_parent_value = self . _parents [ other_parent ] . value <EOL> num_singletons = len ( d_out_d_self . shape ) - len ( parent_shape ) <EOL> if num_singletons > <NUM_LIT:0> : <EOL> extra_singletons = tuple ( range ( num_singletons ) ) <EOL> result = np . sum ( d_out_d_self * other_parent_value , axis = extra_singletons , keepdims = False ) <EOL> else : <EOL> result = d_out_d_self * other_parent_value <EOL> assert len ( result . shape ) == len ( parent_shape ) <EOL> original_singletons = tuple ( np . where ( np . array ( parent_shape ) == <NUM_LIT:1> ) [ <NUM_LIT:0> ] ) <EOL> return np . sum ( result , axis = original_singletons , keepdims = True ) <EOL> class MatDet ( Differentiable ) : <EOL> __slots__ = [ '<STR_LIT:A>' ] <EOL> def __init__ ( self , A , axis = None , keepdims = True ) : <EOL> super ( MatDet , self ) . __init__ ( ( A , ) ) <EOL> self . A = A <EOL> def _compute_value ( self ) : <EOL> return np . linalg . det ( self . A . value ) <EOL> def _local_grad ( self , parent , d_out_d_self ) : <EOL> det = self . _compute_value ( ) <EOL> return d_out_d_self * det * np . linalg . inv ( self . A . value ) . T <EOL> class MatLogDet ( Differentiable ) : <EOL> pass <EOL> class MatTrace ( Differentiable ) : <EOL> pass <EOL> class Transpose ( Differentiable ) : <EOL> __slots__ = [ '<STR_LIT:A>' , '<STR_LIT>' ] <EOL> def __init__ ( self , A , axes = None ) : <EOL> super ( Transpose , self ) . __init__ ( ( A , ) ) <EOL> self . A = A <EOL> self . axes = axes <EOL> def _compute_value ( self ) : <EOL> return np . transpose ( self . A . value , axes = self . axes ) <EOL> def _local_grad ( self , parent , d_out_d_self ) : <EOL> if self . axes is None : <EOL> return np . transpose ( d_out_d_self ) <EOL> else : <EOL> return np . transpose ( d_out_d_self , axes = np . argsort ( self . axes ) ) <EOL> class Reshape ( Differentiable ) : <EOL> __slots__ = [ '<STR_LIT:A>' , '<STR_LIT>' ] <EOL> def __init__ ( self , A , new_shape ) : <EOL> super ( Reshape , self ) . __init__ ( ( A , ) ) <EOL> self . A = A <EOL> self . new_shape = new_shape <EOL> def _compute_value ( self ) : <EOL> return np . reshape ( self . A . value , self . new_shape ) <EOL> def _local_grad ( self , parent , d_out_d_self ) : <EOL> return np . reshape ( d_out_d_self , self . A . shape ) <EOL> class Concatenate ( Differentiable ) : <EOL> __slots__ = [ '<STR_LIT>' ] <EOL> def __init__ ( self , axis , * args ) : <EOL> super ( Concatenate , self ) . __init__ ( args ) <EOL> self . axis = axis <EOL> def _compute_value ( self ) : <EOL> return np . concatenate ( [ p . value for p in self . _parents ] , axis = self . axis ) <EOL> def _local_grad ( self , parent_ix , d_out_d_self ) : <EOL> start_ix = sum ( [ p . shape [ self . axis ] for p in self . _parents [ <NUM_LIT:0> : parent_ix ] ] ) <EOL> end_ix = start_ix + self . _parents [ parent_ix ] . shape [ self . axis ] <EOL> return index_along_axis ( d_out_d_self , self . axis , start_ix , end_ix ) <EOL> class ListToArray ( Differentiable ) : <EOL> """<STR_LIT>""" <EOL> __slots__ = [ ] <EOL> def __init__ ( self , * args ) : <EOL> super ( ListToArray , self ) . __init__ ( args ) <EOL> def _compute_value ( self ) : <EOL> return np . concatenate ( [ p . value [ None , : ] for p in self . _parents ] , axis = <NUM_LIT:0> ) <EOL> def _local_grad ( self , parent_ix , d_out_d_self ) : <EOL> return d_out_d_self [ parent_ix , : ] <EOL> def index_along_axis ( array , axis , start , end ) : <EOL> """<STR_LIT>""" <EOL> full_slice = [ slice ( None ) , ] * array . ndim <EOL> full_slice [ axis ] = slice ( start , end ) <EOL> return array [ full_slice ] <EOL> class TensorMult ( Differentiable ) : <EOL> __slots__ = [ '<STR_LIT>' ] <EOL> def __init__ ( self , A , B , axes ) : <EOL> super ( TensorMult , self ) . __init__ ( ( A , B ) ) <EOL> self . axes = axes <EOL> def _compute_value ( self ) : <EOL> A = self . _parents [ <NUM_LIT:0> ] . value <EOL> B = self . _parents [ <NUM_LIT:1> ] . value <EOL> return np . tensordot ( A , B , self . axes ) <EOL> def _local_grad ( self , parent , d_out_d_self ) : <EOL> diff = lambda A , B : [ a for a in A if a not in B ] <EOL> rank = lambda L : list ( np . argsort ( np . argsort ( L ) ) ) <EOL> val = [ p . value for p in self . _parents ] <EOL> axes = self . axes <EOL> n_axes = len ( axes [ <NUM_LIT:0> ] ) <EOL> ignore_dims = [ diff ( range ( val [ i ] . ndim ) , axes [ i ] ) for i in ( <NUM_LIT:0> , <NUM_LIT:1> ) ] <EOL> ignore_ndims = [ len ( x ) for x in ignore_dims ] <EOL> output_dims = ( range ( ignore_ndims [ <NUM_LIT:0> ] ) , <EOL> range ( ignore_ndims [ <NUM_LIT:0> ] , ignore_ndims [ <NUM_LIT:0> ] + ignore_ndims [ <NUM_LIT:1> ] ) ) <EOL> X , Y = parent , <NUM_LIT:1> - parent <EOL> wrong_order = np . tensordot ( val [ Y ] , d_out_d_self , ( ignore_dims [ Y ] , output_dims [ Y ] ) ) <EOL> permutation = [ None ] * val [ X ] . ndim <EOL> for final , cur in zip ( list ( axes [ X ] ) + ignore_dims [ X ] , <EOL> rank ( axes [ Y ] ) + range ( n_axes , val [ X ] . ndim ) ) : <EOL> permutation [ final ] = cur <EOL> return np . transpose ( wrong_order , permutation ) <EOL> class Identity ( Differentiable ) : <EOL> __slots__ = [ ] <EOL> def __init__ ( self , A ) : <EOL> super ( Identity , self ) . __init__ ( ( A , ) ) <EOL> def _compute_value ( self ) : <EOL> return self . _parents [ <NUM_LIT:0> ] . value <EOL> def _local_grad ( self , parent_ix , d_out_d_self ) : <EOL> return d_out_d_self </s>
