mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Mosaic] Fix assert

Browse Source 
PiperOrigin-RevId: 560756391
This commit is contained in:
Tomás Longeri 2023-08-28 10:46:34 -07:00 committed by jax authors
parent 70206ee6cd
commit c3e624ad8a
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jaxlib/mosaic/python/apply_vector_layout.py
View File

@ -184,7 +184,8 @@ class VectorLayout:
# Tiling should neatly divide the target shape, so that every vector
# register ends up having the same structure.
# Also, every tile should occupy a fixed number of sublanes.
assert self.tiling[0] * self.tiling[1] % self.packing * TARGET_SHAPE[1] == 0
assert ((self.tiling[0] * self.tiling[1]) % (self.packing * TARGET_SHAPE[1])
== 0)
# Offsets should not exceed the tile size. The data always starts within the
# first tile of a vreg.
assert all(0 <= (o or 0) < t for o, t in zip(self.offsets, self.tiling))