Apply Autogression over 3d numpy array and predicting array in Python

I have a 3d Numpy array with observations over 3 points in time:

a = np.array(time,x,y)

My goal is to run an autoregression on the observations over the 3 points in time and predict an array containing the future values.

I tried doing this using a for loop over rows and columns, but then my autoregression only takes the past 3 values (one per dimension) into account leading to the error:

ValueError: maxlag should be < nobs

How can I apply the autoregression over the three dimensions using not only a specific x and y, but all entries (as the trend is assumed to be the same)? Also, is there a faster way then looping as the final np.array will be much bigger?

Thank you!

This is the code I wrote:

a = np.array([[[ 5,  8,  6,  5],
              [ 8,  2,  3,  3],
              [ 0,  4,  7,  1],
              [ 9,  9,  6,  1]],

             [[ 2,  0, 10,  9],
              [ 3,  6,  8,  7],
              [ 1,  4,  2,  9],
              [ 2,  4,  9,  2]],

             [[ 6,  9,  9,  1],
              [ 8,  7,  6,  1],
              [ 5,  3,  2,  2],
              [ 7,  7,  2,  5]]])


# Loop over each index in the array
for row in range(a.shape[0] ) :
    for col in range(a.shape[1] ) :
        data = (a[0,row,col], a[1,row,col], a[2,row,col])

        # fit model
        model = AR(data)
        model_fit = model.fit()

        # make prediction
        yhat = model_fit.predict(len(data), len(data))

        print(yhat)