From a73cace509615b624e78adb5f2d398e662cc102c Mon Sep 17 00:00:00 2001
From: Albert Zeyer <albzey@gmail.com>
Date: Thu, 12 Oct 2023 17:59:47 +0200
Subject: [PATCH] Dim math cache better, use same_base (#1419)

---
 returnn/tensor/_dim_extra.py | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/returnn/tensor/_dim_extra.py b/returnn/tensor/_dim_extra.py
index e773338f4a..a61ba019ce 100644
--- 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)