Matplotlib多个子图共用colorbar

问题

在Matplotlib中，如何使多个子图（subplot）共用一个colorbar？

代码

import matplotlib.pyplot as plt
from matplotlib import colors
import numpy as np

dataset1 = np.random.randint(0, 50, size=(10, 10))
dataset2 = np.random.randint(50, 100, size=(10, 10))

vmin = min(np.min(dataset1), np.min(dataset2))
vmax = max(np.max(dataset1), np.max(dataset2))
norm = colors.Normalize(vmin=vmin, vmax=vmax)

fig = plt.figure()
ax = fig.add_subplot(121)
bx = fig.add_subplot(122)
a = ax.pcolormesh(dataset1, norm=norm, cmap=plt.get_cmap('rainbow'))
b = bx.pcolormesh(dataset2, norm=norm, cmap=plt.get_cmap('rainbow'))
fig.colorbar(a, ax=[ax, bx], shrink=0.5)
# fig.colorbar(a, ax=ax, shrink=0.5)
# fig.colorbar(b, ax=bx, shrink=0.5)
plt.show()
           

效果