Log scale in 3dplot using matploblib

I would like to create a 3D plot with the z axis in log scale. This is an example code.

import numpy as npfrom mpl_toolkits.mplot3d import Axes3Dimport matplotlib.pyplot as pltdef fun(x, y):    return x**2 + yfig = plt.figure()ax = fig.add_subplot(111, projection='3d')x = y = np.arange(-3.0, 3.0, 0.05)X, Y = np.meshgrid(x, y)zs = np.array([fun(x,y) for x,y in zip(np.ravel(X), np.ravel(Y))])Z = zs.reshape(X.shape)ax.plot_surface(X, Y, Z)ax.set_xlabel('X Label')ax.set_ylabel('Y Label')ax.set_zlabel('Z Label')plt.show()

I tried ax.set_zscale('log') and ax.zaxis.set_scale('log') but that is not working. It should be possible though. If I use ax.plot_surface(X, Y, np.log(Z)) the z ticks are off. I could not find an answer to this seemingly simple problem. I hope you can help.