<s> import numpy as np <EOL> import cudarray as ca <EOL> from . . wrap import blas <EOL> from . . linalg import matmul_shape <EOL> class Dot ( object ) : <EOL> def __init__ ( self , a , b , out = None ) : <EOL> self . batch_size = a . shape [ <NUM_LIT:0> ] <EOL> self . a = a <EOL> self . b = b <EOL> if a . dtype != b . dtype : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> out_shape = ( self . batch_size , ) + matmul_shape ( a . shape [ <NUM_LIT:1> : ] , b . shape [ <NUM_LIT:1> : ] ) <EOL> if out is None : <EOL> out = base . empty ( out_shape , dtype = a . dtype ) <EOL> else : <EOL> if out_shape != out . shape : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if a . dtype != out . dtype : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> self . out = out <EOL> a_stride = np . prod ( a . shape [ <NUM_LIT:1> : ] ) <EOL> b_stride = np . prod ( b . shape [ <NUM_LIT:1> : ] ) <EOL> out_stride = np . prod ( out . shape [ <NUM_LIT:1> : ] ) <EOL> self . blas_batch = blas . BLASBatch_f ( <EOL> a . _data , b . _data , out . _data , self . batch_size , a_stride , b_stride , <EOL> out_stride <EOL> ) <EOL> if a . ndim == b . ndim == <NUM_LIT:3> : <EOL> m , k = a . shape [ <NUM_LIT:1> : <NUM_LIT:3> ] <EOL> n = b . shape [ <NUM_LIT:2> ] <EOL> def fun ( ) : <EOL> self . blas_batch . gemm ( blas . no_trans_op , blas . no_trans_op , m , n , <EOL> k , <NUM_LIT:1.0> , <NUM_LIT:0.0> ) <EOL> return self . out <EOL> self . perform = fun <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> def dot ( a , b , out = None ) : <EOL> return Dot ( a , b , out ) . perform ( ) </s>
