JAX scan loop not compiling when having a single tile_map call

[balancap]

There is a weird bug where it seems that the call is forwarded to CPU XLA backend when there is a single tile_map call in a JAX loop.

Minimal reproducer:

from functools import partial

import jax
import numpy as np

from tessellate_ipu import tile_map, tile_put_sharded
from tessellate_ipu.lax import sqrt_inplace_p

data = np.array([1, 2, 3], np.float32)
tiles = (1, 2, 5)

def inner_scan(carry, _):
    return tile_map(jax.lax.sqrt_p, carry), None

@partial(jax.jit, backend="ipu")
def compute_fn(input):
    x = tile_put_sharded(input, tiles)
    return jax.lax.scan(inner_scan, x, None, length=4)

output = compute_fn(data)