#NEWLINE# @ primitive #NEWLINE# def convolve( A , B , axes = None , dot_axes = [ ( ) , ( ) ] , mode = 'full' ) : #NEWLINE# #INDENT# assert ( mode in [ 'valid' , 'full' ] ) , 'Mode {0} not yet implemented' . format ( mode ) #NEWLINE# if ( axes is None ) : #NEWLINE# #INDENT# axes = [ list ( range ( A . ndim ) ) , list ( range ( A . ndim ) ) ] #NEWLINE# #UNINDENT# wrong_order = any ( [ ( B . shape [ ax_B ] < A . shape [ wrong_order ] ) for ( ax_A , ax_B ) in zip ( * axes ) ] ) #NEWLINE# if wrong_order : #NEWLINE# #INDENT# if ( ( mode == 'valid' ) and ( not all ( [ ( B . shape [ ax_B ] <= A . shape [ ax_A ] ) for ( ax_A , ax_B ) in zip ( * axes ) ] ) ) ) : #NEWLINE# #INDENT# raise Exception ( 'One array must be larger than the other along all convolved dimensions' ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# if ( ( mode != 'full' ) or ( B . size <= A . size ) ) : #NEWLINE# #INDENT# i1 = ( ( B . ndim - len ( dot_axes [ 1 ] ) ) - len ( axes [ 1 ] ) ) #NEWLINE# i2 = ( ( ( i1 + A . ndim ) - len ( dot_axes [ 0 ] ) ) - len ( axes [ 0 ] ) ) #NEWLINE# i3 = ( i2 + len ( axes [ 0 ] ) ) #NEWLINE# ignore_B = list ( range ( i1 ) ) #NEWLINE# ignore_A = list ( range ( i1 , i2 ) ) #NEWLINE# conv = list ( range ( i2 , i3 ) ) #NEWLINE# return convolve ( B , A , axes = axes [ : : ( - 1 ) ] , dot_axes = dot_axes [ : : ( - 1 ) ] , mode = mode ) . transpose ( ( ( ignore_A + ignore_B ) + conv ) ) #NEWLINE# #UNINDENT# #UNINDENT# #UNINDENT# if ( mode == 'full' ) : #NEWLINE# #INDENT# B = pad_to_full ( B , A , axes [ : : ( - 1 ) ] ) #NEWLINE# #UNINDENT# B_view_shape = list ( B . shape ) #NEWLINE# B_view_strides = list ( B . strides ) #NEWLINE# flipped_idxs = ( [ slice ( None ) ] * A . ndim ) #NEWLINE# for ( ax_A , ax_B ) in zip ( * axes ) : #NEWLINE# #INDENT# B_view_shape . append ( ( abs ( ( B . shape [ ax_B ] - A . shape [ ax_A ] ) ) + 1 ) ) #NEWLINE# B_view_strides . append ( B . strides [ ax_B ] ) #NEWLINE# B_view_shape [ ax_B ] = A . shape [ ax_A ] #NEWLINE# flipped_idxs [ ax_A ] = slice ( None , None , ( - 1 ) ) #NEWLINE# #UNINDENT# B_view = as_strided ( B , B_view_shape , B_view_strides ) #NEWLINE# A_view = A [ flipped_idxs ] #NEWLINE# all_axes = [ ( list ( axes [ i ] ) + list ( dot_axes [ i ] ) ) for i in [ 0 , 1 ] ] #NEWLINE# return einsum_tensordot ( A_view , B_view , all_axes )