Fix quantizer base dtype

PiperOrigin-RevId: 660442990
google · Aug 7, 2024 · f99d63c · f99d63c
1 parent 81154d8
commit f99d63c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/praxis/layers/video/quantizer.py b/praxis/layers/video/quantizer.py
@@ -153,7 +153,7 @@ def __call__(self, inputs: JTensor) -> tuple[JTensor, NestedMap]:
     result_dict = NestedMap()
     group_size = int(math.ceil(self.embedding_dim / self.num_token_groups))
     base = jnp.power(
-        2, jnp.arange(self.embedding_dim, dtype=jnp.uint8) % group_size
+        2, jnp.arange(self.embedding_dim, dtype=jnp.uint32) % group_size
     )
     samples = inputs >= 0
     quantized = jnp.where(samples, 1.0, -1.0)
@@ -215,7 +215,7 @@ def __call__(self, inputs: JTensor) -> tuple[JTensor, NestedMap]:
   def decode_ids(self, inputs: JTensor) -> JTensor:
     group_size = int(math.ceil(self.embedding_dim / self.num_token_groups))
     base = jnp.power(
-        2, jnp.arange(self.embedding_dim, dtype=jnp.uint8) % group_size
+        2, jnp.arange(self.embedding_dim, dtype=jnp.uint32) % group_size
     )
     if self.num_token_groups == 1:
       inputs = inputs[..., None]