How to convert 3D histogram to heatmap

I am wondering if there is a way to map a 3D histogram from matplotlib to a color heatmap? For example I have a 3D histogram as

How can I convert it to a heatmap where the height is shown as color of the heatmap?

In Matplotlib there is also Axes.hist2d or pyplot.hist2d.

import matplotlib.pyplot as plt
import numpy as np
from numpy.random import default_rng

rng = default_rng(0)

x, y = rng.uniform(0, 4, size=(2, 100))

fig, ax = plt.subplots()
ax.hist2d(x, y, bins=4, range=[[0, 4], [0, 4]], cmap='Blues')
ax.set_aspect('equal')

plt.show()

To add a colorbar to estimate the counts in each bin you can use the following code. I've also demonstrated how to obtain the resultant hist matrix that matplotlib generates.

import matplotlib.pyplot as plt
import numpy as np
from numpy.random import default_rng

rng = default_rng(0)

x, y = rng.uniform(0, 4, size=(2, 100))

fig, ax = plt.subplots()
hist, xedges, yedges, im = ax.hist2d(
    x, y, bins=4, range=[[0, 4], [0, 4]], cmap='Blues'
)
ax.set_aspect('equal')
fig.colorbar(im, ax=ax, label='count')

print(hist)
# [[ 5.  7.  6.  4.]
#  [ 5.  6.  4.  7.]
#  [ 5.  8.  6.  6.]
#  [ 7.  1.  9. 14.]]

plt.show()