def stacked_matmul(a, b):
 if (a.shape[1] > 50):
  out = np.empty_like(a)
  for i in range(a.shape[0]):
   out[i] = np.dot(a[i], b[i])
  return out
 else:
  return np.einsum('...ij,...jk->...ik', a, b)

def stacked_matmul ( a , b ) : if ( a . shape [ 1 ] > 50 ) : out = np . empty_like ( a ) for i in range ( a . shape [ 0 ] ) : out [ i ] = np . dot ( a [ i ] , b [ i ] ) return out else : return np . einsum ( '...ij,...jk->...ik' , a , b )
stacked matrix multiply: out[i .