Colorbar for inset in python

I am trying to add an additional small colorbar for the inset axis. The current code, without that, is

 import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib import cm
 import matplotlib.colors as colors
 from mpl_toolkits.axes_grid1.inset_locator import mark_inset
 A = np.linspace(1,20,20)
 B = A
 X,Y = np.meshgrid(A,B)
 Z = X**2 + Y**2    
 fig, ax = plt.subplots()
 im = ax.pcolor(X, Y, Z, cmap='hot_r') 
 ax.set_xlabel('x',fontsize=labelsize)
 ax.set_ylabel('y',fontsize=labelsize)
 ca = fig.colorbar(im)#, shrink=0.5, aspect=5)
 axins = ax.inset_axes([0.1, 0.5, 0.25, 0.25])
 axins.pcolor(A[0:4], B[0:4], Z[0:4,0:4], cmap='hot_r')
 axins.tick_params(axis='both', which='major', labelsize=11)
 for axis in ['top','bottom','left','right']:
     axins.spines[axis].set_linewidth(1)
     axins.spines[axis].set_color('gray')
 mark_inset(ax, axins, loc1=2, loc2=4, fc="none", ec='gray', lw=1)
 plt.tight_layout()

You could create an additional inset axis for the colorbar. E.g. located just right of the inset. Then create a colorbar proving this axis (cax=...).

Please note that pcolor creates faces (large pixels) between the given x and y positions. So, you need one row and one column more of position then the number of colors. The current version of matplotlib gives a warning in case too many colors (or not enough x and y positions) are given.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import mark_inset

A = np.linspace(1, 20, 20)
B = A
X, Y = np.meshgrid(A, B)
Z = X ** 2 + Y ** 2
fig, ax = plt.subplots()
im = ax.pcolor(X, Y, Z[:-1, :-1], cmap='hot_r')
ax.set_xlabel('x', fontsize=12)
ax.set_ylabel('y', fontsize=12)
ca = fig.colorbar(im)  # , shrink=0.5, aspect=5)
axins = ax.inset_axes([0.1, 0.5, 0.25, 0.25])
axins_cbar = ax.inset_axes([0.37, 0.5, 0.02, 0.25])
img_in = axins.pcolor(A[0:5], B[0:5], Z[0:4, 0:4], cmap='hot_r')
axins.tick_params(axis='both', which='major', labelsize=11)
for axis in ['top', 'bottom', 'left', 'right']:
    axins.spines[axis].set_linewidth(1)
    axins.spines[axis].set_color('gray')
mark_inset(ax, axins, loc1=2, loc2=4, fc="none", ec='gray', lw=1)
fig.colorbar(img_in, cax=axins_cbar)
plt.tight_layout()

plt.show()