Replace tuple unpacking with jnp.s_ for variable indexing

PiperOrigin-RevId: 668199363
google · Aug 28, 2024 · b391500 · b391500
1 parent 8bbf1f5
commit b391500
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/aqt/__init__.py b/aqt/__init__.py
@@ -13,4 +13,4 @@
 # limitations under the License.
 """Accurate Quantized Training library."""
 
-__version__ = "0.8.0"
+__version__ = "0.8.1"
diff --git a/aqt/jax/v2/tiled_dot_general.py b/aqt/jax/v2/tiled_dot_general.py
@@ -245,7 +245,7 @@ def unapply(self, tiled_x: jnp.ndarray) -> jnp.ndarray:
     num_bcast_axes = len(self.get_broadcasted_tile_map_indexes())
     # All elements of broadcast axes are the same, thus we take the first one.
     first_index = (0,) * num_bcast_axes
-    x = tiled_x[*first_index]
+    x = tiled_x[jnp.s_[first_index]]
     x = x.reshape(self.untiled_shape)
     return x