Problem regarding axis shape in 3d surface plot in python

This is my first attempt at 3d surface plotting. Spent hours on trying to get a plot to show but no luck yet. Here is my script to do a surface plot of z_ax with respect to x_ax and y_ax. The z_ax variable is generated from some other script, but i thats irrelevant, so i just copy pasted the array just for the sake of reproduction.

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.axes3d import Axes3D
import numpy as np
z_ax = np.array([[98.0952381 , 98.37627192, 98.22459584, 97.96470716, 97.97332409,
        98.55658199, 97.59356635, 97.72653459],
       [98.78787879, 99.56700585, 99.26385681, 99.41539461, 99.43703447,
        99.37211316, 99.28240025, 99.30172134],
       [99.43722944, 99.35050877, 99.35046189, 99.3287864 , 99.50632254,
        99.3432448 , 99.39993814, 99.38291653],
       [99.13419913, 99.52370643, 99.52367206, 99.44787269, 99.49766153,
        99.51645497, 99.55459326, 99.51824185],
       [99.48051948, 99.45875731, 99.58140878, 99.58861102, 99.50632254,
        99.4154157 , 99.44324157, 99.51282884],
       [99.61038961, 99.41545789, 99.53810624, 99.51282884, 99.4630175 ,
        99.54532333, 99.51128982, 99.55613294],
       [99.43722944, 99.52370643, 99.4948037 , 99.58861102, 99.51498354,
        99.48036952, 99.51747603, 99.5398939 ],
       [99.48051948, 99.58865555, 99.50923788, 99.49117679, 99.41971245,
        99.52367206, 99.55459326, 99.5398939 ],
       [99.65367965, 99.48040701, 99.66801386, 99.577785  , 99.44569548,
        99.55254042, 99.57315187, 99.59402403],
       [99.39393939, 99.67525438, 99.48036952, 99.59943705, 99.5756106 ,
        99.52367206, 99.52366223, 99.51282884]])

x_ax = np.linspace(0.1,0.8,8) # the test size
y_ax = np.linspace(0.01,1,10) # zero class weight 

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.plot_surface(x_ax, y_ax, z_ax)
plt.show()

However, when i run the script, i get:

Traceback (most recent call last):
File "tmp.py", line 30, in <module>
ax.plot_surface(x_ax, y_ax, z_ax)
File "C:\Python27\lib\site-packages\mpl_toolkits\mplot3d\axes3d.py", line 1567, in plot_surface
X, Y, Z = np.broadcast_arrays(X, Y, Z)
File "C:\Python27\lib\site-packages\numpy\lib\stride_tricks.py", line 249, 
in broadcast_arrays
shape = _broadcast_shape(*args)
File "C:\Python27\lib\site-packages\numpy\lib\stride_tricks.py", line 184, 
in _broadcast_shape
b = np.broadcast(*args[:32])
ValueError: shape mismatch: objects cannot be broadcast to a single shape

I checked the shapes of the co-ordinates, they seem to match for the surface plot. So, i dont understand what i am getting wrong. Any help is highly appreciated.