Keep square pixels + ability to have any rectangular shape zoom with Plotly.js, with a plot with 2 layers

In Zoom on a Plotly heatmap, the accepted answer provides a way to zoom on a heatmap that:

• has squared pixels (required)
• allows any rectangular-shaped zoom (and not stuck on the heatmap's aspect ratio).

This stops working if we have a Plotly.js plot with 2 layers : 1 "heatmap" and 1 "image": in the following snippet, the zooming feature is stuck on the heatmap's original aspect ratio, which I don't want:

How to allow free rectangular shape zoom instead?

const z = Array.from({ length: 50 }, () => Array.from({ length: 20 }, () => Math.floor(Math.random() * 255)));
const z2 = [[[255, 0, 0], [0, 255, 255]], [[0, 0, 255], [255, 0, 255]]];
const data = [{ type: "image", z: z2 }, { type: "heatmap", z: z, opacity: 0.3 }];
const layout = { xaxis: { scaleanchor: "y", scaleratio: 1 } };
Plotly.newPlot("plot", data, layout).then(afterPlot);
function afterPlot(gd) {
    const xrange = gd._fullLayout.xaxis.range;
    const yrange = gd._fullLayout.yaxis.range;
    const xrange_init = [...xrange];
    const yrange_init = [...yrange];
    const zw0 = xrange[1] - xrange[0];
    const zh0 = yrange[1] - yrange[0];
    const r0 = Number((zw0 / zh0).toPrecision(6));
    const update = { "xaxis.range": xrange, "yaxis.range": yrange, "xaxis.scaleanchor": false };
    Plotly.relayout(gd, update);
    gd.on("plotly_relayout", relayoutHandler);
    function relayoutHandler(e) {
        if (e.width || e.height) {
            return unbindAndReset(gd, relayoutHandler);
        }
        if (e["xaxis.autorange"] || e["yaxis.autorange"]) {
            [xrange[0], xrange[1]] = xrange_init;
            [yrange[0], yrange[1]] = yrange_init;
            return Plotly.relayout(gd, update);
        }
        const zw1 = xrange[1] - xrange[0];          
        const zh1 = yrange[1] - yrange[0];
        const r1 = Number((zw1 / zh1).toPrecision(6));
        if (r1 === r0) {
            return;
        }
        const [xmin, xmax] = getExtremes(gd, 0, "x");
        const [ymin, ymax] = getExtremes(gd, 0, "y");
        if (r1 > r0) {
            const extra = ((zh1 * r1) / r0 - zh1) / 2;
            expandAxisRange(yrange, extra, ymin, ymax);
        }
        if (r1 < r0) {
            const extra = ((zw1 * r0) / r1 - zw1) / 2;
            expandAxisRange(xrange, extra, xmin, xmax);
        }
        Plotly.relayout(gd, update);
    }
}
function unbindAndReset(gd, handler) {
    gd.removeListener("plotly_relayout", handler); 
    return Plotly.relayout(gd, { xaxis: { scaleanchor: "y", scaleratio: 1, autorange: true }, yaxis: { autorange: true } }).then(afterPlot);
}
function getExtremes(gd, traceIndex, axisId) {
    const extremes = gd._fullData[traceIndex]._extremes[axisId];
    return [extremes.min[0].val, extremes.max[0].val];
}
function expandAxisRange(range, extra, min, max) {
    let shift = 0;
    if (range[0] - extra < min) {
        const out = min - (range[0] - extra);
        const room = max - (range[1] + extra);
        shift = out <= room ? out : (out + room) / 2;
    } else if (range[1] + extra > max) {
        const out = range[1] + extra - max;
        const room = range[0] - extra - min;
        shift = out <= room ? -out : -(out + room) / 2;
    }
    range[0] = range[0] - extra + shift;
    range[1] = range[1] + extra + shift;
}

<script src="https://cdn.plot.ly/plotly-2.22.0.min.js"></script>
<div id="plot"></div>

from Keep square pixels + ability to have any rectangular shape zoom with Plotly.js, with a plot with 2 layers