jax

be7fe878 - [pallas:triton] Elide `program_id` calls where launch grid dimension is 1.

[pallas:triton] Elide `program_id` calls where launch grid dimension is 1. This may allow for parts of indexing calculations to be optimized away. PiperOrigin-RevId: 678321871