autoset modality_num_dim if modality_default_shape specified

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "transfusion-pytorch"
-version = "0.10.1"
+version = "0.10.2"
 description = "Transfusion in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_transfusion.py

@@ -234,8 +234,7 @@ def test_velocity_consistency():
         num_text_tokens = 12,
         dim_latent = 384,
         channel_first_latent = True,
-        modality_default_shape = ((4, 4)),
-        modality_num_dim = 2,
+        modality_default_shape = (4, 4),
         modality_encoder = mock_encoder,
         modality_decoder = mock_decoder,
         transformer = dict(

transfusion_pytorch/transfusion.py

@@ -1276,11 +1276,30 @@ def __init__(
 
         self.to_modality_shape_fn = cast_tuple(to_modality_shape_fn, self.num_modalities)
 
+        # default token lengths for respective modality
+        # fallback if the language model does not come up with valid dimensions
+
+        if not exists(modality_default_shape) or is_bearable(modality_default_shape, tuple[int, ...]):
+            modality_default_shape = (modality_default_shape,) * self.num_modalities
+
+        self.modality_default_shape = modality_default_shape
+
+        assert len(self.modality_default_shape) == self.num_modalities
+
+        self.fallback_to_default_shape_if_invalid = fallback_to_default_shape_if_invalid
+
+        # default `modality_num_dim` to `len(modality_default_shape)` if latter is specified but former not
+
+        modality_num_dim = default(modality_num_dim, tuple(len(shape) for shape in self.modality_default_shape))
+
         # specifying the number of dimensions for the modality, which will be hard validated
 
         self.modality_num_dim = cast_tuple(modality_num_dim, self.num_modalities)
+
         assert len(self.modality_num_dim) == self.num_modalities
 
+        assert all([not exists(ndim) or not exists(shape) or len(shape) == ndim for ndim, shape in zip(self.modality_num_dim, self.modality_default_shape)])
+
         # whether to add an extra axial positional embedding per modality
 
         self.add_pos_emb = cast_tuple(add_pos_emb, self.num_modalities)
@@ -1318,18 +1337,6 @@ def __init__(
 
         self.maybe_add_temp_batch_dim = add_temp_batch_dim if modality_encoder_decoder_requires_batch_dim else identity
 
-        # default token lengths for respective modality
-        # fallback if the language model does not come up with valid dimensions
-
-        if not exists(modality_default_shape) or is_bearable(modality_default_shape, tuple[int, ...]):
-            modality_default_shape = (modality_default_shape,) * self.num_modalities
-
-        self.modality_default_shape = modality_default_shape
-
-        assert len(self.modality_default_shape) == self.num_modalities
-
-        self.fallback_to_default_shape_if_invalid = fallback_to_default_shape_if_invalid
-
         # store number of text tokens
 
         self.num_text_tokens = num_text_tokens