Update with min and with_max to handle change from ragged to fixed

NVIDIA-Merlin · Mar 8, 2023 · 5494898 · 5494898
1 parent b329001
commit 5494898
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/merlin/dtypes/shape.py b/merlin/dtypes/shape.py
@@ -187,10 +187,16 @@ def with_dim(self, index, value):
         return replace(self, dims=tuple(new_dims))
 
     def with_dim_min(self, index, value):
-        return self.with_dim(index, self.dims[index].with_min(value))
+        new_dim = self.dims[index].with_min(value)
+        if new_dim.is_uniform:
+            new_dim = Dimension(value)
+        return self.with_dim(index, new_dim)
 
     def with_dim_max(self, index, value):
-        return self.with_dim(index, self.dims[index].with_max(value))
+        new_dim = self.dims[index].with_max(value)
+        if new_dim.is_uniform:
+            new_dim = Dimension(value)
+        return self.with_dim(index, new_dim)
 
     @property
     def min(self) -> Tuple: