I am afraid I need more information to try and help. Axes could refer to either the matplotlib object (subplot below) or to the dimensions in a numpy array.
Is the issue with the number of subplots generated and their layout? For example, getting a single row with 4 subplots instead of a 2x2 grid of subplots.
Or is the issue with the shape of the array returned as ax? For example, getting a 1d array of length 4 (thus, shape (4,)) for a plot with 1 row and 4 columns instead of one with shape (1,4)
Or maybe something else entirely?