Computing Matrix in 4d is not Working

Hello,

I have these array:

A_ax(d,d,b,m) 
B_by(d,d,b,m)

When I have A_ax(2,2,2,2)the code is working and I am getting the correct result:

d=size(A_ax,1);
m=size(A_ax,3);

P = zeros(d,d,m,m);

for x=1:m
    for y=1:m
        for a=1:d
            for b=1:d
                P(a,b,x,y) = real( trace(kron(A_ax(:,:,x,a),B_by(:,:,y,b))*rho));
            end
        end
    end
end

However when I have that A_ax(4,4,2,2) code is not working. This is because a is going untill 4 but I dont have 4. I know the error but I could not find a correction for that:

How can I solve the problem?