Skip to content

Commit

Permalink
[pallas:mosaic_gpu] Removed a stale TODO
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 680931423
  • Loading branch information
superbobry authored and Google-ML-Automation committed Oct 1, 2024
1 parent f62941d commit 14ef2b6
Showing 1 changed file with 0 additions and 3 deletions.
3 changes: 0 additions & 3 deletions jax/_src/pallas/mosaic_gpu/lowering.py
Original file line number Diff line number Diff line change
Expand Up @@ -457,9 +457,6 @@ def store(
# We have to do some work to make sure that consecutive stores are not
# going to be writing to the same location, or else we'll end up with
# multiple concurrent writes and a racy program.
# TODO(apaszke,slebedev): In most cases output index maps depend only on
# parallel grid axes and in that case we can simply move the store to
# happen after the loop.
# TODO(apaszke,slebedev): This still diverges significantly from the TPU
# semantics in that it will move on to the next SMEM output slice even if
# it's not storing the previous one.
Expand Down

0 comments on commit 14ef2b6

Please sign in to comment.