account for indices being optional in fsq for images

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.14.29"
+version = "1.14.30"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

vector_quantize_pytorch/finite_scalar_quantization.py

@@ -4,6 +4,7 @@
 """
 
 from __future__ import annotations
+from functools import wraps
 from typing import List, Tuple
 
 import torch
@@ -25,6 +26,14 @@ def default(*args):
             return arg
     return None
 
+def maybe(fn):
+    @wraps(fn)
+    def inner(x, *args, **kwargs):
+        if not exists(x):
+            return x
+        return fn(x, *args, **kwargs)
+    return inner
+
 def pack_one(t, pattern):
     return pack([t], pattern)
 
@@ -179,7 +188,11 @@ def forward(self, z):
             z = z.float()
 
         codes = self.quantize(z)
+
+        # returning indices could be optional
+
         indices = None
+
         if self.return_indices:
             indices = self.codes_to_indices(codes)
 
@@ -200,10 +213,10 @@ def forward(self, z):
             out = unpack_one(out, ps, 'b * d')
             out = rearrange(out, 'b ... d -> b d ...')
 
-            indices = unpack_one(indices, ps, 'b * c')
+            indices = maybe(unpack_one)(indices, ps, 'b * c')
 
         if not self.keep_num_codebooks_dim and self.return_indices:
-            indices = rearrange(indices, '... 1 -> ...')
+            indices = maybe(rearrange)(indices, '... 1 -> ...')
 
         # return quantized output and indices