x-transformers 1.27.20__tar.gz → 1.27.22__tar.gz

{x_transformers-1.27.20/x_transformers.egg-info → x_transformers-1.27.22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.27.20
+Version: 1.27.22
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.27.20 → x_transformers-1.27.22}/README.md

@@ -1437,11 +1437,10 @@ model = XValAutoregressiveWrapper(model)
 
 ids = torch.randint(0, 4, (1, 777))
 nums = torch.randn(1, 777)
-mask = torch.ones(1, 777).bool()
 
 # train on a lot of data above
 
-loss = model(ids, nums, mask = mask)
+loss = model(ids, nums)
 loss.backward()
 
 # then generate

{x_transformers-1.27.20 → x_transformers-1.27.22}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '1.27.20',
+  version = '1.27.22',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

{x_transformers-1.27.20 → x_transformers-1.27.22}/x_transformers/xval.py

@@ -260,10 +260,19 @@ class XValAutoregressiveWrapper(nn.Module):
         inp, target = x[:, :-1], x[:, 1:]
         x_num_inp, x_num_target = x_num[:, :-1], x_num[:, 1:]
 
+        # ignore index
+
+        target_mask = target != self.ignore_index
+
+        # key padding mask
+
         mask = kwargs.get('mask', None)
-        if exists(mask) and mask.shape[1] == x.shape[1]:
-            mask = mask[:, :-1]
-            kwargs['mask'] = mask
+        if exists(mask):
+            target_mask &= mask
+
+            if mask.shape[1] == x.shape[1]:
+                mask = mask[:, :-1]
+                kwargs['mask'] = mask
 
         logits, numerical_pred = self.net(inp, x_num_inp, **kwargs)
 
@@ -271,17 +280,23 @@ class XValAutoregressiveWrapper(nn.Module):
 
         cross_entropy_loss = F.cross_entropy(logits, target, reduction = 'none', ignore_index = self.ignore_index)
 
-        target_mask = target != self.ignore_index
+        # protect against nan in `x_num` input tensor
+
+        target_is_number_mask = target == self.net.numerical_token_id
+        x_num_target = x_num_target.masked_fill(~target_is_number_mask, 0.)
+
+        # numerical mse loss
 
         numerical_mse_loss = F.mse_loss(numerical_pred, x_num_target, reduction = 'none')
 
         numerical_mse_loss = numerical_mse_loss * target_mask
+        numerical_mse_loss = numerical_mse_loss.masked_fill(~target_is_number_mask, 0.)
 
-        loss = cross_entropy_loss + numerical_mse_loss * self.numerical_loss_weight
+        # combine losses
 
-        if exists(mask):
-            loss = loss[mask]
+        loss = cross_entropy_loss + numerical_mse_loss * self.numerical_loss_weight
 
+        loss = loss[target_mask]
         loss = loss.mean()
 
         if not return_loss_breakdown:

{x_transformers-1.27.20 → x_transformers-1.27.22/x_transformers.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.27.20
+Version: 1.27.22
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang