Is scan supported in pallas?

[hr0nix]

I have a kernel code that contains jax.lax.map. It runs fine with interpret=True, however lowering to triton fails with the following error:

E         jax_triton.pallas.triton_lowering.TritonLoweringException: Exception while lowering eqn:
E           a:i32[6] = scan[
E           jaxpr={ lambda ; b:Ref{int32[384]} c:i32[]. let
E               d:i32[] = mul c 64
E               e:i32[64] = iota[dimension=0 dtype=int32 shape=(64,)]
E               f:i32[64] = add e d
E               g:bool[64] = lt f 384
E               h:i32[64] <- b[d:d+64]
E               i:i32[] = reduce_min[axes=(0,)] h
E             in (i,) }
E           length=6
E           linear=(False, False)
E           num_carry=0
E           num_consts=1
E           reverse=False
E           unroll=1
E         ] j k

Is it because scan is not supported or is there some other problem? Happy to provide more details if necessary.

[hr0nix]

On a side note, I can't help but notice that pl.load mask (g in the code above) isn't being used in the code, is that expected?

[hr0nix]

On a side note, I can't help but notice that pl.load mask (g in the code above) isn't being used in the code, is that expected?

It appears that masked reads don't work correctly with interpret=True, which is likely related to this.

[sharadmv]

I suspect you're running into an unsupported use-case for scan. Scan automatically slices into its inputs and outputs which isn't supported in Triton afaik. Could you post the full traceback so I could double check it?

Re:masking could you open that as a separate issue with a repro?

[hr0nix]

Full stacktrace: https://gist.github.com/hr0nix/195f1ece2e6cde792cd0ae0e2fbf6357

After carefully looking at it, it indeed looks like I hit a non-implemented or non-supported code path.

[sharadmv]

Yes I can confirm that is the issue:

>     if num_extensive: raise NotImplementedError
E     NotImplementedError