Dim math cache better, use same_base (#1419)

rwth-i6 · Oct 12, 2023 · a73cace · a73cace
1 parent dcd082b
commit a73cace
Showing 1 changed file with 11 additions and 11 deletions.
diff --git a/returnn/tensor/_dim_extra.py b/returnn/tensor/_dim_extra.py
@@ -1995,7 +1995,7 @@ def __add__(self: Dim, other):
         :rtype: Dim
         """
         cache_key = ("add", other)
-        cache = self._make_extra().cache_dim_math
+        cache = self.get_same_base()._make_extra().cache_dim_math
         if cache_key in cache:
             return cache[cache_key]
         term = _OpLinearTerm.from_dim(self)
@@ -2011,7 +2011,7 @@ def __radd__(self: Dim, other):
         :rtype: Dim
         """
         cache_key = ("add_left", other)
-        cache = self._make_extra().cache_dim_math
+        cache = self.get_same_base()._make_extra().cache_dim_math
         if cache_key in cache:
             return cache[cache_key]
         term = _OpLinearTerm.from_dim(self)
@@ -2034,7 +2034,7 @@ def sub_right(self: Dim, other):
         :rtype: Dim
         """
         cache_key = ("sub", other)
-        cache = self._make_extra().cache_dim_math
+        cache = self.get_same_base()._make_extra().cache_dim_math
         if cache_key in cache:
             return cache[cache_key]
         term = _OpLinearTerm.from_dim(self)
@@ -2050,7 +2050,7 @@ def sub_left(self: Dim, other):
         :rtype: Dim
         """
         cache_key = ("sub_left", other)
-        cache = self._make_extra().cache_dim_math
+        cache = self.get_same_base()._make_extra().cache_dim_math
         if cache_key in cache:
             return cache[cache_key]
         term = _OpLinearTerm.from_dim(self)
@@ -2065,7 +2065,7 @@ def __mul__(self: Dim, other):
         :rtype: Dim
         """
         cache_key = ("mul", other)
-        cache = self._make_extra().cache_dim_math
+        cache = self.get_same_base()._make_extra().cache_dim_math
         if cache_key in cache:
             return cache[cache_key]
         term = _OpLinearTerm.from_dim(self)
@@ -2080,7 +2080,7 @@ def __rmul__(self: Dim, other):
         :rtype: Dim
         """
         cache_key = ("mul_left", other)
-        cache = self._make_extra().cache_dim_math
+        cache = self.get_same_base()._make_extra().cache_dim_math
         if cache_key in cache:
             return cache[cache_key]
         term = _OpLinearTerm.from_dim(self)
@@ -2095,7 +2095,7 @@ def __floordiv__(self: Dim, other):
         :rtype: Dim
         """
         cache_key = ("floordiv", other)
-        cache = self._make_extra().cache_dim_math
+        cache = self.get_same_base()._make_extra().cache_dim_math
         if cache_key in cache:
             return cache[cache_key]
         term = _OpLinearTerm.from_dim(self)
@@ -2117,7 +2117,7 @@ def div_left(self: Dim, other):
         :rtype: Dim
         """
         cache_key = ("truediv_left", other)
-        cache = self._make_extra().cache_dim_math
+        cache = self.get_same_base()._make_extra().cache_dim_math
         if cache_key in cache:
             return cache[cache_key]
         term = _OpLinearTerm.from_dim(self)
@@ -2132,7 +2132,7 @@ def div_right(self: Dim, other):
         :rtype: Dim
         """
         cache_key = ("truediv", other)
-        cache = self._make_extra().cache_dim_math
+        cache = self.get_same_base()._make_extra().cache_dim_math
         if cache_key in cache:
             return cache[cache_key]
         term = _OpLinearTerm.from_dim(self)
@@ -2147,7 +2147,7 @@ def ceildiv_left(self: Dim, other):
         :rtype: Dim
         """
         cache_key = ("ceildiv_left", other)
-        cache = self._make_extra().cache_dim_math
+        cache = self.get_same_base()._make_extra().cache_dim_math
         if cache_key in cache:
             return cache[cache_key]
         term = _OpLinearTerm.from_dim(self)
@@ -2162,7 +2162,7 @@ def ceildiv_right(self: Dim, other):
         :rtype: Dim
         """
         cache_key = ("ceildiv", other)
-        cache = self._make_extra().cache_dim_math
+        cache = self.get_same_base()._make_extra().cache_dim_math
         if cache_key in cache:
             return cache[cache_key]
         term = _OpLinearTerm.from_dim(self)