Shape mismatch when using lstsq_simulate with TF backend

[CKrawczyk]

I came across your paper on the arxiv and wanted to test out some of the examples. I ran into an issue when running lstsq_simulate with the TF backend:

ValueError: Shape must be rank 4 but is rank 3 for '{{node concat}} = ConcatV2[N=2, T=DT_FLOAT, Tidx=DT_INT32](zeros, PartitionedCall_1, concat/axis)' with input shapes: [0,60,60,1], [60,60,1], [].

Looking around the source code I think I tracked down the issue to these lines:
https://github.com/giga-lens/gigalens/blob/master/src/gigalens/tf/simulator.py#L119-L121

Specifically, there is a shape mismatch between img and tmp.

I took a look at some of the other branches and this seems to be fixed on the tf-demo branch:
https://github.com/giga-lens/gigalens/blob/tf-demo-update/gigalens/tf/simulator.py#L119-L123

But I am reluctant to switch to that branch as it seems to be out of date with the master branch. I figured I would flag it up as this bug appears to be present in the package's current pypi release.

[CKrawczyk]

I just tested the JAX backend and found the same issue there as well:

TypeError: Cannot concatenate arrays with different numbers of dimensions: got (0, 120, 120, 1), (120, 120, 1).

Looks to be these lines: https://github.com/giga-lens/gigalens/blob/master/src/gigalens/jax/simulator.py#L110-L114