In lax.scan, let unroll=0 unroll the whole loop

[carlosgmartin]

It would be convenient to have an integer value for the unroll argument of lax.scan that unrolls the whole loop. (As opposed to the unergonomic and confusing situation of having to switch between integers and the boolean True, not to be confused with 1, which is confusingly equivalent to False.)

I argue that 0 would be a natural choice for such a value.

unroll=n splits a sequence of size k into chunks of size n followed by a chunk of size k mod n.

Thus unroll=0 should split a sequence of size k into chunks of size 0 (i.e. nothing) followed by a chunk of size k mod 0.

But k mod 0 == k. (See, e.g., here, here, here, and here.)

Therefore, unroll=0 should split a sequence into nothing followed by a chunk of size k (i.e., the whole sequence).

If this proposal is agreeable, I can create a short PR for this.

[mattjj]

SGTM! Thanks for the clear explanation, and for teaching me about k mod 0 == k.

Would we keep the current behavior on bools, or would you propose deprecating that?

[carlosgmartin]

@mattjj Personally, I'd be in favor of deprecating bool inputs, since this proposal would render it obsolete.

[mattjj]

I think the main question is how much churn it'd cause (in breaking downstream user code, and updating human brains), and how the cost/benefit of that churn balances out. My guess is the cost/benefit would be kind of high and thus it would be a low priority cleanup at best.

It's not even totally clear to me how desirable it is.... getting full unrolling by passing unroll=True seems slightly easier to guess than unroll=0. But I don't have a strong feeling about that.

Having unroll=int_0_or_1 meaning the opposite of unroll=bool(int_0_or_1), or equivalently having unroll=int(some_bool) meaning the opposite of unroll=int(some_bool), is indeed a bit weird but my a priori hunch is that it's unlikely to cause issues because people are unlikely to forget whether they're passing an int or a bool.

So how about we proceed on supporting unroll=0 while keeping unroll: int | bool, but we could open a separate issue for deprecating bool-typed unroll values?

[carlosgmartin]

That sounds good. We can go with just adding support for unroll=0 for now.

I've created a PR at #29630.