BUG: TesselateIPU IO tiles demo

[balancap]

Description

TessellateIPU IO tiles demo (https://github.com/graphcore-research/tessellate-ipu/pull/20) fails when the following line is removed:

iobuffer_init = jax.device_put(iobuffer_init, jax.devices("ipu")[0])

Two things then happen:

• 2 IPU PjRt executables are created, instead of one. Leading to IPU:0 reset + reload + run phase;
• iobuffer is not properly initialized in the second reset + run => using default zero values; Should raise an error, or best, trigger the proper HOST buffer transfer;

What jax/jaxlib version are you using?

No response

Which accelerator(s) are you using?

IPU

Additional System Info

No response