python - xor matrix multiplication for AES mix column stage -

hi i'm writing program aes mix column stage. here have multiply 2 matrices of (4,4) shape. difference while multiplying 2 matrices have take 'xor' instead of have add. e.g

a = np.array([[1,2],[3,4]]) b = np.array([[5,6],[7,8]]) np.dot(a,b) # gives [[(1*5+2*7),(1*6+2*8)][(3*5+4*7),(3*6+4*8)]]             # want [[((1*5)^(2*7)),((1*6)^(2*8))][((3*5)^(4*7)),((3*6)^(4*8))]] 

here's solution loops

result = [[0,0,0,0],           [0,0,0,0],           [0,0,0,0],           [0,0,0,0]]  # iterate through rows of x in range(len(x)):    # iterate through columns of y    j in range(len(y[0])):       # iterate through rows of y       k in range(len(y)):          result[i][j] = result[i][j] ^ (x[i][k] * y[k][j]) 

how achieve without using loops?

xor_ab=np.bitwise_xor.reduce(a[...,none]*b,axis=1) 

for explanation, consider rectangular problem easier identification:

a=np.arange(12).reshape(4,3).astype(object) b=np.arange(12).reshape(3,4).astype(object) 

object provide python int arbitrary precision aes.

products obtained broadcasting,

c=a[...,none]*b   # dims :  (4,3,1) * ((1),3,4) -> (4,3,4)  , c_ijk =a_ij*b_jk 

the dot product obtained :

dot_ab=c.sum(axis=1)  #  ->(4,4) in [734]: (dot_ab==a.dot(b)).all() out[734]: true  

then change equivalent xor function :

xor_ab=np.bitwise_xor.reduce(a[...,none]*b,axis=1) 

as alternative, can interpret loops numba (0.23):

from numba import jit @jit(nopython=true) def xor(x,y):     result=np.zeros((4,4),np.uint64)     in range(len(x)):     # iterate through columns of y       j in range(y.shape[1]):         # iterate through rows of y         k in range(len(y)):             result[i,j] = result[i,j] ^ (x[i,k] * y[k,j])     return result 

for impressive efficiency gain,due optimal memory usage. limited 32 bits , b:

in [790]: %timeit xor(a,b) 1000000 loops, best of 3: 580 ns per loop  in [791]: %timeit xor_ab=np.bitwise_xor.reduce(a[...,none]*b,axis=1) 100000 loops, best of 3: 13.2 µs per loop  in [792] (xor(a,b)==np.bitwise_xor.reduce(a[...,none]*b,axis=1)).all() out[792]: true