[Pallas] Reject 64-bit input tensors and fix tiling ZeroDivisionError

[norx1991]

Summary

TPU does not natively support 64-bit element types. This PR:

1. Runtime (runtime/__init__.py): _pallas_check_dtypes raises TypeError if any tensor arg uses int64, uint64, or float64. Runs before the cache check in all three launchers so it cannot be bypassed.

2. Block spec tiling (backend.py): Cap bitwidth to 32 in _get_pallas_required_alignment to prevent ZeroDivisionError when 64-bit dtypes appear at compile time (128 * (32 // 64) = 0).

3. LONG_INT_TYPE (_testing.py): New constant (int32 on Pallas, int64 elsewhere) for examples/tests that use long integers but don't need 64-bit range. Updated cross_entropy example and test to use it.

Note: internal dtype promotion (e.g. torch.sum on int32 → int64, argmax → int64) is handled by JAX, which silently narrows without jax_enable_x64.

[AmesingFlank]

The changes lgtm, although I'm not sure if we really wanna do this. One could argue that if the user wrote int64/float64, we're expected to uphold that contract. Might be better to just throw an error, explicitly calling out that 64-bit types are not supported?

[norx1991]

The changes lgtm, although I'm not sure if we really wanna do this. One could argue that if the user wrote int64/float64, we're expected to uphold that contract. Might be better to just throw an error, explicitly calling out that 64-bit types are not supported?

Thanks for the quick reply. Several kernels use x64 (e.g., cross_entropy, as listed in #2010), so throwing errors mean we would have to either create an almost-duplicated example, or decide not to support it. Or do you think we should guard this under a flag?

[AmesingFlank]

The changes lgtm, although I'm not sure if we really wanna do this. One could argue that if the user wrote int64/float64, we're expected to uphold that contract. Might be better to just throw an error, explicitly calling out that 64-bit types are not supported?

Thanks for the quick reply. Several kernels use x64 (e.g., cross_entropy, as listed in #2010), so throwing errors mean we would have to either create an almost-duplicated example, or decide not to support it. Or do you think we should guard this under a flag?

I personally think we should just use a mechanism similar to this one for these tests/examples. E.g. we can do LONG_INT_TYPE and set them to int64 for others and int32 for pallas.

My mental model is that if someone requested int64/float64 dtypes, they must've done so for a reason, and they believe that they need the extra bitwidth. It's likely for these use cases, falling back to int32 is not really useful.

[norx1991]

The changes lgtm, although I'm not sure if we really wanna do this. One could argue that if the user wrote int64/float64, we're expected to uphold that contract. Might be better to just throw an error, explicitly calling out that 64-bit types are not supported?

Thanks for the quick reply. Several kernels use x64 (e.g., cross_entropy, as listed in #2010), so throwing errors mean we would have to either create an almost-duplicated example, or decide not to support it. Or do you think we should guard this under a flag?

I personally think we should just use a mechanism similar to this one for these tests/examples. E.g. we can do LONG_INT_TYPE and set them to int64 for others and int32 for pallas.

My mental model is that if someone requested int64/float64 dtypes, they must've done so for a reason, and they believe that they need the extra bitwidth. It's likely for these use cases, falling back to int32 is not really useful.

@AmesingFlank Thanks, this is a good point. I checked cross_entropy, and there doesn’t seem to be any real need for long there; it looks like it was mainly mirroring the standard PyTorch API. For cases like that, using LONG_INT_TYPE seems reasonable.

For kernels where the data range or precision actually matters, I agree we should raise an explicit error instead of silently narrowing.

So overall I agree that defining LONG_INT_TYPE is the better step forward here: it keeps the choice explicit in the example/test code and has a much smaller blast radius than changing runtime behavior.

I will update this PR to reflect this.

[norx1991]

The changes lgtm, although I'm not sure if we really wanna do this. One could argue that if the user wrote int64/float64, we're expected to uphold that contract. Might be better to just throw an error, explicitly calling out that 64-bit types are not supported?

Thanks for the quick reply. Several kernels use x64 (e.g., cross_entropy, as listed in #2010), so throwing errors mean we would have to either create an almost-duplicated example, or decide not to support it. Or do you think we should guard this under a flag?

I personally think we should just use a mechanism similar to this one for these tests/examples. E.g. we can do LONG_INT_TYPE and set them to int64 for others and int32 for pallas.
My mental model is that if someone requested int64/float64 dtypes, they must've done so for a reason, and they believe that they need the extra bitwidth. It's likely for these use cases, falling back to int32 is not really useful.

@AmesingFlank Thanks, this is a good point. I checked cross_entropy, and there doesn’t seem to be any real need for long there; it looks like it was mainly mirroring the standard PyTorch API. For cases like that, using LONG_INT_TYPE seems reasonable.

For kernels where the data range or precision actually matters, I agree we should raise an explicit error instead of silently narrowing.

So overall I agree that defining LONG_INT_TYPE is the better step forward here: it keeps the choice explicit in the example/test code and has a much smaller blast radius than changing runtime behavior.

Updated.