How to generate all possible 3x3 matrices with fixed first row and column with entries from 0,1,2

Question

I want to fix the first row and first column of the matrices to be 0 1 2..only such matrices is to be filtered among all possible 3x3 matrices with entries from 0,1,2. Later I have to check one condition for those filtered ones. The condition is as follows:

i + j = a_ij (the corresponding entry of one matrix) = m , m + k = a_mk 
j + k = a_jk (entry from j th row and k th column) = n, i + n = a_in 

I have to display only those matrices among all combinations which satisfy a_mk = a_in. Its more like checking associativity. Can anyone help to complete this?

from itertools import product
import numpy as np
m=3
n=3

x = [[list(i[x:x+m]) for x in range(0, len(i), m)] for i in product("012", repeat=m*n)]
print(len(x))
A=[]
for j in x:
  j = np.array(j)
  for i in range(3):
    if j[0][i] == i:
      if j[i][0] == i:
        A.append(j)

print(len(A))

So far I have been able to write so far.

Answer 1

Here is the Modifyed Code

from itertools import product
import numpy as np

m = 3
n = 3

x = [[list(i[x:x+m]) for x in range(0, len(i), m)] for i in product("012", repeat=m*n)]

A = []
for j in x:
    j = np.array(j, dtype=int)

    # Check conditions
    condition_1 = np.all(j[0, :] == np.arange(m))
    condition_2 = np.all(j[:, 0] == np.arange(n))
    condition_3 = np.all(j[1:, :] + np.arange(m) == j[:-1, :])
    condition_4 = np.all(j[:, 1:] + np.arange(n) == j[:, :-1])
    condition_5 = np.all(j[0, :] + np.arange(n) == j[:, 0])
    
    # Check associativity condition
    associativity_condition = np.all(j[1:, 1:] == j[:-1, :-1])
    
    if condition_1 and condition_2 and condition_3 and condition_4 and condition_5 and associativity_condition:
        A.append(j)

print(len(A))
print(A)