returnn 1.20231018.70522__tar.gz → 1.20231018.94805__tar.gz

{returnn-1.20231018.70522/returnn.egg-info → returnn-1.20231018.94805}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20231018.70522
+Version: 1.20231018.94805
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20231018.94805/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20231018.094805'
+long_version = '1.20231018.094805+git.74e0e94'

{returnn-1.20231018.70522 → returnn-1.20231018.94805}/returnn/frontend/_backend.py

@@ -791,11 +791,13 @@ class Backend(Generic[T]):
             a,
             b,
             name=kind,
-            res_dtype="bool",
+            copy_sparse_dim=False,
             allow_broadcast_all_sources=allow_broadcast_all_sources,
             dim_order=dim_order,
         )
-        out.raw_tensor = cls.compare_raw(a_raw, kind, b_raw)
+        out_raw = cls.compare_raw(a_raw, kind, b_raw)
+        out.dtype = cls.get_dtype_name_raw(out_raw)
+        out.raw_tensor = out_raw
         return out
 
     @classmethod
@@ -816,11 +818,12 @@ class Backend(Generic[T]):
             a,
             b,
             name=kind,
-            res_dtype=None,
             allow_broadcast_all_sources=allow_broadcast_all_sources,
             dim_order=dim_order,
         )
-        out.raw_tensor = cls.combine_raw(a_raw, kind, b_raw)
+        out_raw = cls.combine_raw(a_raw, kind, b_raw)
+        out.dtype = cls.get_dtype_name_raw(out_raw)
+        out.raw_tensor = out_raw
         return out
 
     @staticmethod

{returnn-1.20231018.70522 → returnn-1.20231018.94805}/returnn/frontend/_native/module.cpp

@@ -377,7 +377,7 @@ bool PyModuleState::_torchTensorDTypesInit() {
     PyObjectScopedRef mod = PyImport_ImportModule("torch");
     if(!mod) return false;
 
-    int i = 0;
+    unsigned int i = 0;
     #define AddDType(dtype_) \
         assert(i < sizeof(_torchTensorDTypes)/sizeof(_torchTensorDTypes[0])); \
         _torchTensorDTypes[i].dtype = PyObject_GetAttrString(mod, #dtype_); \

{returnn-1.20231018.70522 → returnn-1.20231018.94805}/returnn/frontend/_native/module.hpp

@@ -112,17 +112,17 @@ public:
 
     int pyTraverse(visitproc visit, void *arg) {
         Py_VISIT(_notSpecified);
-        for(int i = 0; i < sizeof(_rawTensorTypes)/sizeof(_rawTensorTypes[0]); ++i)
+        for(unsigned int i = 0; i < sizeof(_rawTensorTypes)/sizeof(_rawTensorTypes[0]); ++i)
             Py_VISIT(_rawTensorTypes[i]);
         Py_VISIT(_tensorType);
         Py_VISIT(_dimType);
         Py_VISIT(_globalBackend);
         Py_VISIT(_backendTensorTypeDispatchTable);
-        for(int i = 0; i < NumBackendsWithCachedOps * NumTOps; ++i)
+        for(unsigned int i = 0; i < NumBackendsWithCachedOps * NumTOps; ++i)
             Py_VISIT(_cachedOps[i]);
         Py_VISIT(_torchTensorType);
         Py_VISIT(_torchBackend);
-        for(int i = 0; i < sizeof(_torchTensorDTypes)/sizeof(_torchTensorDTypes[0]); ++i) {
+        for(unsigned int i = 0; i < sizeof(_torchTensorDTypes)/sizeof(_torchTensorDTypes[0]); ++i) {
             Py_VISIT(_torchTensorDTypes[i].dtype);
             Py_VISIT(_torchTensorDTypes[i].name);
         }
@@ -133,17 +133,17 @@ public:
     int pyClear() {
         Py_CLEAR(_notSpecified);
         _rawTensorTypesLen = 0;
-        for(int i = 0; i < sizeof(_rawTensorTypes)/sizeof(_rawTensorTypes[0]); ++i)
+        for(unsigned int i = 0; i < sizeof(_rawTensorTypes)/sizeof(_rawTensorTypes[0]); ++i)
             Py_CLEAR(_rawTensorTypes[i]);
         Py_CLEAR(_tensorType);
         Py_CLEAR(_dimType);
         Py_CLEAR(_globalBackend);
         Py_CLEAR(_backendTensorTypeDispatchTable);
-        for(int i = 0; i < NumBackendsWithCachedOps * NumTOps; ++i)
+        for(unsigned int i = 0; i < NumBackendsWithCachedOps * NumTOps; ++i)
             Py_CLEAR(_cachedOps[i]);
         Py_CLEAR(_torchTensorType);
         Py_CLEAR(_torchBackend);
-        for(int i = 0; i < sizeof(_torchTensorDTypes)/sizeof(_torchTensorDTypes[0]); ++i) {
+        for(unsigned int i = 0; i < sizeof(_torchTensorDTypes)/sizeof(_torchTensorDTypes[0]); ++i) {
             Py_CLEAR(_torchTensorDTypes[i].dtype);
             Py_CLEAR(_torchTensorDTypes[i].name);
         }
@@ -176,7 +176,7 @@ public:
                 return NULL;
             }
         }
-        for(int i = 0; i < sizeof(_torchTensorDTypes)/sizeof(_torchTensorDTypes[0]); ++i) {
+        for(unsigned int i = 0; i < sizeof(_torchTensorDTypes)/sizeof(_torchTensorDTypes[0]); ++i) {
             if(_torchTensorDTypes[i].dtype == dtype)
                 return _torchTensorDTypes[i].name;
         }

{returnn-1.20231018.70522 → returnn-1.20231018.94805}/returnn/frontend/_native/tensor_ops.cpp

@@ -18,8 +18,8 @@ PyObject* tensorCopy(
 {
     PyObjectScopedRef rawTensor = PyObject_GetAttrString(tensor, "_raw_tensor");
     if(rawTensor == Py_None)
-        return tensorCopyTemplate(modState, tensor, name, NULL);
-    PyObjectScopedRef res = tensorCopyTemplate(modState, tensor, name, NULL);
+        return tensorCopyTemplate(modState, tensor, name);
+    PyObjectScopedRef res = tensorCopyTemplate(modState, tensor, name);
     if(!res) return NULL;
     if(PyObject_SetAttrString(res, "_raw_tensor", rawTensor) < 0) return NULL;
     return res.release();
@@ -64,7 +64,7 @@ PyObject* tensorCopyTemplate(
     PyModuleState* modState,
     PyObject* tensor,
     const char* name,
-    const char* dtype)
+    PyObject* dtype)
 {
     PyObjectScopedRef version = PyObject_GetAttrString(tensor, "version");
     if(!version) return NULL;
@@ -99,8 +99,10 @@ PyObject* tensorCopyTemplate(
         if(!name_) return NULL;
         if(PyDict_SetItemString(kwargs, "name", name_) < 0) return NULL;
     }
-    {
-        PyObjectScopedRef dtype_ = dtype ? PyUnicode_FromString(dtype) : PyObject_GetAttrString(tensor, "dtype");
+    if(dtype && dtype != Py_None) {
+        if(PyDict_SetItemString(kwargs, "dtype", dtype) < 0) return NULL;
+    } else {
+        PyObjectScopedRef dtype_ = PyObject_GetAttrString(tensor, "dtype");
         if(!dtype_) return NULL;
         if(PyDict_SetItemString(kwargs, "dtype", dtype_) < 0) return NULL;
     }
@@ -147,17 +149,22 @@ PyObject* tensorCopyTemplateSimple(
     PyModuleState* modState,
     PyObject* tensor,
     const char* name_,
-    const char* dtype_)
+    PyObject* dtype,
+    bool copySparseDim)
 {
     PyObjectScopedRef name = name_ ? PyUnicode_FromString(name_) : PyObject_GetAttrString(tensor, "name");
     if(!name) return NULL;
-    PyObjectScopedRef dtype = dtype_ ? PyUnicode_FromString(dtype_) : PyObject_GetAttrString(tensor, "dtype");
-    if(!dtype) return NULL;
+    PyObjectScopedRef dtype_;
+    if(!dtype) {
+        dtype_ = PyObject_GetAttrString(tensor, "dtype");
+        if(!dtype_) return NULL;
+        dtype = dtype_.get();
+    }
     PyObjectScopedRef dims = PyObject_GetAttrString(tensor, "_dims");
     if(!dims) return NULL;
 
     PyObjectScopedRef res = PyObject_CallFunctionObjArgs(
-        modState->tensorType(), name.get(), dims.get(), dtype.get(), NULL);
+        modState->tensorType(), name.get(), dims.get(), dtype, NULL);
     if(!res) return NULL;
 
     {
@@ -167,7 +174,7 @@ PyObject* tensorCopyTemplateSimple(
             if(PyObject_SetAttrString(res, "_feature_dim_axis", feature_dim_axis) < 0)
                 return NULL;
     }
-    if(!dtype_) {
+    if(copySparseDim) {
         PyObjectScopedRef sparse_dim = PyObject_GetAttrString(tensor, "sparse_dim");
         if(!sparse_dim) return NULL;
         if(sparse_dim != Py_None)
@@ -300,8 +307,8 @@ PyObject* pyTensorCopyTemplate(PyObject *self, PyObject *args, PyObject *kwargs)
     static const char *kwlist[] = { "tensor", "name", "dtype", NULL };
     PyObject* tensor;
     const char* name = NULL;
-    const char* dtype = NULL;
-    if(!PyArg_ParseTupleAndKeywords(args, kwargs, "O|z$z:tensor_copy_template",
+    PyObject* dtype = NULL;
+    if(!PyArg_ParseTupleAndKeywords(args, kwargs, "O|z$O:tensor_copy_template",
             (char**) kwlist, &tensor, &name, &dtype))
         return NULL;
 
@@ -389,9 +396,9 @@ static bool _isMatchingDimTagsAndRawShape(PyObject* dimTags, PyObject* rawShape,
 }
 
 static bool _checkTensorRawTensorAssignForBackendWithCachedOps(
-    PyModuleState* modState, BackendWithCachedOps backendId, const char* funcName, PyObject* tensor, PyObject* rawTensor
+    PyModuleState* modState, BackendWithCachedOps backendId, const char* funcName, PyObject* tensor, PyObject* rawTensor, bool checkDtype = true
 ) {
-    {
+    if(checkDtype) {
         PyObject* getDTypeOp = modState->cachedOp(TOp_GetDType, backendId);
         if(!getDTypeOp) return false;
         PyObjectScopedRef dtype = PyObject_GetAttrString(tensor, "dtype");
@@ -528,8 +535,8 @@ static bool _getPermutationSupersetToSubset(const char* funcName, ASeqT subset,
             ++count;
         }
         else if(candidates.size() > 1) {
-            size_t maxMatchPriorityIdx;
-            long maxMatchPriority;
+            size_t maxMatchPriorityIdx = 0;
+            long maxMatchPriority = -1;
             int countSameMatchPriority = 0;
             for(size_t j = 0; j < candidates.size(); ++j) {
                 PyObject* dim_ = subset.getItem(candidates[j]);
@@ -667,7 +674,7 @@ template<bool bIsSubset>
 static PyObject* _compareOrCombine_subsetDims(
     PyModuleState* modState,
     const char* rawOpName, bool resultIsBool,
-    PyObject* permuteOp, PyObject* reshapeOp, PyObject* getShapeOp, PyObject* rawOp,
+    PyObject* permuteOp, PyObject* reshapeOp, PyObject* getShapeOp, PyObject* getDtypeOp, PyObject* rawOp,
     PyObject* a, PyObject* b,
     PyObject* aRawTensor, PyObject* bRawTensor,
     PyTupleOrListStaticRef<true> aDims, PyTupleOrListStaticRef<true> bDims,
@@ -681,12 +688,14 @@ static PyObject* _compareOrCombine_subsetDims(
         rawTensorExt = _permuteAndExtend(rawOpName, permuteOp, reshapeOp, getShapeOp, a, aDims, aRawTensor, bDims, outPermutation);
 
     // Now create the result.
-    PyObjectScopedRef res = tensorCopyTemplateSimple(modState, bIsSubset ? a : b, rawOpName, resultIsBool ? "bool" : NULL);
-    if(!res) return NULL;
     PyObjectScopedRef resRawTensor = PyObject_CallFunctionObjArgs(
         rawOp, bIsSubset ? aRawTensor : rawTensorExt.get(), bIsSubset ? rawTensorExt.get() : bRawTensor, NULL);
     if(!resRawTensor) return NULL;
-    if(PyObject_SetAttrString(res, "raw_tensor", resRawTensor) < 0) return NULL;
+    PyObjectScopedRef dtype = PyObject_CallFunctionObjArgs(getDtypeOp, resRawTensor.get(), NULL);
+    if(!dtype) return NULL;
+    PyObjectScopedRef res = tensorCopyTemplateSimple(modState, bIsSubset ? a : b, rawOpName, dtype, !resultIsBool);
+    if(!res) return NULL;
+    if(PyObject_SetAttrString(res, "_raw_tensor", resRawTensor) < 0) return NULL;
     return res.release();
 }
 
@@ -956,7 +965,7 @@ static PyObject* tensorCopyCompatibleToDims(const char* funcName, PyModuleState*
     }
 
     if(outRawTensor)
-        if(PyObject_SetAttrString(outTensor, "raw_tensor", outRawTensor) < 0)
+        if(PyObject_SetAttrString(outTensor, "_raw_tensor", outRawTensor) < 0)
             return NULL;
 
     if(versionInt == 1) {
@@ -1010,12 +1019,12 @@ static PyObject* compareOrCombine(
     bool resultIsBool,
     PyModuleState* modState,
     const char* rawOpName,
-    PyObject* rawOp, PyObject* permuteOp, PyObject* reshapeOp, PyObject* getShapeOp, PyObject* convertToTensorLikeOp,
+    PyObject* rawOp, PyObject* permuteOp, PyObject* reshapeOp, PyObject* getShapeOp, PyObject* getDtypeOp, PyObject* convertToTensorLikeOp,
     bool needConvertToTensor,
     bool allowBroadcastAllSources,
     PyObject* dimOrder
 ) {
-    if(!rawOp || !permuteOp || !reshapeOp || !getShapeOp || !convertToTensorLikeOp) return NULL;
+    if(!rawOp || !permuteOp || !reshapeOp || !getShapeOp || !getDtypeOp || !convertToTensorLikeOp) return NULL;
 
     {
         int a_is_tensor = PyObject_IsInstance(a, modState->tensorType());
@@ -1032,8 +1041,6 @@ static PyObject* compareOrCombine(
                 return NULL;
             }
             // assume the non-Tensor obj is is scalar
-            PyObjectScopedRef res = tensorCopyTemplateSimple(modState, a_is_tensor ? a : b, rawOpName, resultIsBool ? "bool" : NULL);
-            if(!res) return NULL;
             PyObjectScopedRef aRawTensor, bRawTensor;
             if(a_is_tensor) {
                 assert(a_is_tensor && !b_is_tensor);
@@ -1056,7 +1063,11 @@ static PyObject* compareOrCombine(
             PyObjectScopedRef resRawTensor = PyObject_CallFunctionObjArgs(
                 rawOp, aRawTensor ? aRawTensor.get() : a, bRawTensor.get() ? bRawTensor.get() : b, NULL);
             if(!resRawTensor) return NULL;
-            if(PyObject_SetAttrString(res, "raw_tensor", resRawTensor) < 0) return NULL;
+            PyObjectScopedRef dtype = PyObject_CallFunctionObjArgs(getDtypeOp, resRawTensor.get(), NULL);
+            if(!dtype) return NULL;
+            PyObjectScopedRef res = tensorCopyTemplateSimple(modState, a_is_tensor ? a : b, rawOpName, dtype, !resultIsBool);
+            if(!res) return NULL;
+            if(PyObject_SetAttrString(res, "_raw_tensor", resRawTensor) < 0) return NULL;
             return res.release();
         }
         if(!a_is_tensor && !b_is_tensor) {
@@ -1066,31 +1077,6 @@ static PyObject* compareOrCombine(
         // both are Tensor
     }
 
-    if(!resultIsBool) {
-        PyObjectScopedRef aDtype = PyObject_GetAttrString(a, "dtype");
-        if(!aDtype) return NULL;
-        if(!PyUnicode_Check(aDtype)) {
-            PyErr_Format(
-                PyExc_TypeError,
-                "compareOrCombine: a.dtype did not return a string, from dtype %R", aDtype.get());
-            return NULL;
-        }
-        PyObjectScopedRef bDtype = PyObject_GetAttrString(b, "dtype");
-        if(!bDtype) return NULL;
-        if(!PyUnicode_Check(bDtype)) {
-            PyErr_Format(
-                PyExc_TypeError,
-                "compareOrCombine: b.dtype did not return a string, from dtype %R", bDtype.get());
-            return NULL;
-        }
-        if(PyUnicode_Compare(aDtype, bDtype) != 0) {
-            PyErr_Format(
-                PyExc_ValueError,
-                "compareOrCombine: a.dtype != b.dtype, from a.dtype %R and b.dtype %R", aDtype.get(), bDtype.get());
-            return NULL;
-        }
-    }
-
     PyObjectScopedRef aDims = PyObject_GetAttrString(a, "_dims");
     if(!aDims) return NULL;
     if(!PyTuple_Check(aDims)) {
@@ -1125,31 +1111,37 @@ static PyObject* compareOrCombine(
 
     // first very fast path check, check exact identity of dims
     if(_isSameSeqFast(aDimsSeq, bDimsSeq) && (dimOrder == Py_None || _isSameSeqFast(aDimsSeq, dimOrderSeq))) {
-        PyObjectScopedRef res = tensorCopyTemplateSimple(modState, a, rawOpName, resultIsBool ? "bool" : NULL);
-        if(!res) return NULL;
         PyObjectScopedRef resRawTensor = PyObject_CallFunctionObjArgs(rawOp, aRawTensor.get(), bRawTensor.get(), NULL);
         if(!resRawTensor) return NULL;
-        if(PyObject_SetAttrString(res, "raw_tensor", resRawTensor) < 0) return NULL;
+        PyObjectScopedRef dtype = PyObject_CallFunctionObjArgs(getDtypeOp, resRawTensor.get(), NULL);
+        if(!dtype) return NULL;
+        PyObjectScopedRef res = tensorCopyTemplateSimple(modState, a, rawOpName, dtype, !resultIsBool);
+        if(!res) return NULL;
+        if(PyObject_SetAttrString(res, "_raw_tensor", resRawTensor) < 0) return NULL;
         return res.release();
     }
 
     // check b is scalar
     if(bDimsSeq.size() == 0 && (dimOrder == Py_None || _isSameSeqFast(aDimsSeq, dimOrderSeq))) {
-        PyObjectScopedRef res = tensorCopyTemplateSimple(modState, a, rawOpName, resultIsBool ? "bool" : NULL);
-        if(!res) return NULL;
         PyObjectScopedRef resRawTensor = PyObject_CallFunctionObjArgs(rawOp, aRawTensor.get(), bRawTensor.get(), NULL);
         if(!resRawTensor) return NULL;
-        if(PyObject_SetAttrString(res, "raw_tensor", resRawTensor) < 0) return NULL;
+        PyObjectScopedRef dtype = PyObject_CallFunctionObjArgs(getDtypeOp, resRawTensor.get(), NULL);
+        if(!dtype) return NULL;
+        PyObjectScopedRef res = tensorCopyTemplateSimple(modState, a, rawOpName, dtype, !resultIsBool);
+        if(!res) return NULL;
+        if(PyObject_SetAttrString(res, "_raw_tensor", resRawTensor) < 0) return NULL;
         return res.release();
     }
 
     // check a is scalar
     if(aDimsSeq.size() == 0 && (dimOrder == Py_None || _isSameSeqFast(bDimsSeq, dimOrderSeq))) {
-        PyObjectScopedRef res = tensorCopyTemplateSimple(modState, b, rawOpName, resultIsBool ? "bool" : NULL);
-        if(!res) return NULL;
         PyObjectScopedRef resRawTensor = PyObject_CallFunctionObjArgs(rawOp, aRawTensor.get(), bRawTensor.get(), NULL);
         if(!resRawTensor) return NULL;
-        if(PyObject_SetAttrString(res, "raw_tensor", resRawTensor) < 0) return NULL;
+        PyObjectScopedRef dtype = PyObject_CallFunctionObjArgs(getDtypeOp, resRawTensor.get(), NULL);
+        if(!dtype) return NULL;
+        PyObjectScopedRef res = tensorCopyTemplateSimple(modState, b, rawOpName, dtype, !resultIsBool);
+        if(!res) return NULL;
+        if(PyObject_SetAttrString(res, "_raw_tensor", resRawTensor) < 0) return NULL;
         return res.release();
     }
 
@@ -1159,7 +1151,7 @@ static PyObject* compareOrCombine(
         if(_isSeqSubsetFast(bDimsSeq, aDimsSeq, outPermutation) && (dimOrder == Py_None || _isSameSeqFast(aDimsSeq, dimOrderSeq)))
             return _compareOrCombine_subsetDims<true>(
                 modState, rawOpName, resultIsBool,
-                permuteOp, reshapeOp, getShapeOp, rawOp,
+                permuteOp, reshapeOp, getShapeOp, getDtypeOp, rawOp,
                 a, b,
                 aRawTensor, bRawTensor,
                 aDimsSeq, bDimsSeq,
@@ -1172,7 +1164,7 @@ static PyObject* compareOrCombine(
         if(_isSeqSubsetFast(aDimsSeq, bDimsSeq, outPermutation) && (dimOrder == Py_None || _isSameSeqFast(bDimsSeq, dimOrderSeq)))
             return _compareOrCombine_subsetDims<false>(
                 modState, rawOpName, resultIsBool,
-                permuteOp, reshapeOp, getShapeOp, rawOp,
+                permuteOp, reshapeOp, getShapeOp, getDtypeOp, rawOp,
                 a, b,
                 aRawTensor, bRawTensor,
                 aDimsSeq, bDimsSeq,
@@ -1185,7 +1177,7 @@ static PyObject* compareOrCombine(
         if(_isSeqSubsetReorderFast(bDimsSeq, aDimsSeq, outPermutation) && (dimOrder == Py_None || _isSameSeqFast(aDimsSeq, dimOrderSeq)))
             return _compareOrCombine_subsetDims<true>(
                 modState, rawOpName, resultIsBool,
-                permuteOp, reshapeOp, getShapeOp, rawOp,
+                permuteOp, reshapeOp, getShapeOp, getDtypeOp, rawOp,
                 a, b,
                 aRawTensor, bRawTensor,
                 aDimsSeq, bDimsSeq,
@@ -1198,7 +1190,7 @@ static PyObject* compareOrCombine(
         if(_isSeqSubsetReorderFast(aDimsSeq, bDimsSeq, outPermutation) && (dimOrder == Py_None || _isSameSeqFast(bDimsSeq, dimOrderSeq)))
             return _compareOrCombine_subsetDims<false>(
                 modState, rawOpName, resultIsBool,
-                permuteOp, reshapeOp, getShapeOp, rawOp,
+                permuteOp, reshapeOp, getShapeOp, getDtypeOp, rawOp,
                 a, b,
                 aRawTensor, bRawTensor,
                 aDimsSeq, bDimsSeq,
@@ -1306,17 +1298,7 @@ static PyObject* compareOrCombine(
                 PyList_SET_ITEM(allDims.get(), i, outDims[i]);
         }
 
-        PyObjectScopedRef res;
-        {
-            PyObjectScopedRef name = PyUnicode_FromString(rawOpName);
-            if(!name) return NULL;
-            PyObjectScopedRef dtype = resultIsBool ? PyUnicode_InternFromString("bool") : PyObject_GetAttrString(a, "dtype");
-            if(!dtype) return NULL;
-            res = PyObject_CallFunctionObjArgs(
-                modState->tensorType(), name.get(), allDims.get(), dtype.get(), NULL);
-            if(!res) return NULL;
-        }
-
+        PyObjectScopedRef resRawTensor;
         {
             std::vector<int> outPermutation;
             PyObjectScopedRef aRawTensorExt = _permuteAndExtend(
@@ -1330,10 +1312,21 @@ static PyObject* compareOrCombine(
                 b, bDimsSeq, bRawTensor,
                 allDimsSeq, outPermutation);
             if(!bRawTensorExt) return NULL;
-            PyObjectScopedRef resRawTensor = PyObject_CallFunctionObjArgs(
+            resRawTensor = PyObject_CallFunctionObjArgs(
                 rawOp, aRawTensorExt.get(), bRawTensorExt.get(), NULL);
             if(!resRawTensor) return NULL;
-            if(PyObject_SetAttrString(res, "raw_tensor", resRawTensor) < 0) return NULL;
+        }
+
+        PyObjectScopedRef res;
+        {
+            PyObjectScopedRef name = PyUnicode_FromString(rawOpName);
+            if(!name) return NULL;
+            PyObjectScopedRef dtype = PyObject_CallFunctionObjArgs(getDtypeOp, resRawTensor.get(), NULL);
+            if(!dtype) return NULL;
+            res = PyObject_CallFunctionObjArgs(
+                modState->tensorType(), name.get(), allDims.get(), dtype.get(), NULL);
+            if(!res) return NULL;
+            if(PyObject_SetAttrString(res, "_raw_tensor", resRawTensor) < 0) return NULL;
         }
 
         {
@@ -1389,6 +1382,7 @@ static PyObject* compareOrCombineViaCached(
         modState->cachedOp(TOp_Permute, backendId),
         modState->cachedOp(TOp_Reshape, backendId),
         modState->cachedOp(TOp_GetShape, backendId),
+        modState->cachedOp(TOp_GetDType, backendId),
         modState->cachedOp(TOp_ConvertToTensorLike, backendId),
         needConvertToTensor,
         allowBroadcastAllSources,
@@ -1701,31 +1695,14 @@ static PyObject* _tensorUnaryFunc(PyModuleState* modState, PyObject* tensor) {
         if(!func) return NULL;
         PyObjectScopedRef resRawTensor = PyObject_CallFunctionObjArgs(func, rawTensor.get(), NULL);
         if(!resRawTensor) return NULL;
-        PyObjectScopedRef dtype;
-        const char* dtypeStr = NULL;
-        if(op == TOp_Abs) {
-            /*
-            In case of abs() on a complex tensor, the result is a float tensor.
-            In principle, the logic should be like::
-                if in_dtype.startswith("complex"):
-                    num_bits = int(out.dtype[len("complex") :])
-                    out_dtype = f"float{num_bits // 2}"
-            For simplicity, we just take over whatever dtype the result has.
-            */
-            PyObject* getDtypeOp = modState->cachedOp(TOp_GetDType, BWCO_Torch);
-            if(!getDtypeOp) return NULL;
-            dtype = PyObject_CallFunctionObjArgs(getDtypeOp, resRawTensor.get(), NULL);
-            if(!dtype) return NULL;
-            if(!PyUnicode_Check(dtype)) {
-                PyErr_Format(PyExc_TypeError, "tensor_abs: expected dtype to be str, got %R", dtype.get());
-                return NULL;
-            }
-            dtypeStr = PyUnicode_AsUTF8(dtype);
-            if(!dtypeStr) return NULL;
-        }
-        PyObjectScopedRef res = tensorCopyTemplateSimple(modState, tensor, rawOpName(op), dtypeStr);
+        PyObject* getDtypeOp = modState->cachedOp(TOp_GetDType, BWCO_Torch);
+        if(!getDtypeOp) return NULL;
+        // Just overtake the result dtype. In case of abs(), it might change, but maybe also in other cases.
+        PyObjectScopedRef dtype = PyObject_CallFunctionObjArgs(getDtypeOp, resRawTensor.get(), NULL);
+        if(!dtype) return NULL;
+        PyObjectScopedRef res = tensorCopyTemplateSimple(modState, tensor, rawOpName(op), dtype);
         if(!res) return NULL;
-        if(!_checkTensorRawTensorAssignForBackendWithCachedOps(modState, BWCO_Torch, rawOpName(op), res, resRawTensor))
+        if(!_checkTensorRawTensorAssignForBackendWithCachedOps(modState, BWCO_Torch, rawOpName(op), res, resRawTensor, false))
             return NULL;
         if(PyObject_SetAttrString(res, "_raw_tensor", resRawTensor) < 0) return NULL;
         return res.release();

{returnn-1.20231018.70522 → returnn-1.20231018.94805}/returnn/frontend/_native/tensor_ops.hpp

@@ -8,8 +8,9 @@ class PyModuleState;
 // generic
 
 PyObject* tensorCopy(PyModuleState* modState, PyObject* tensor, const char* name = NULL);
-PyObject* tensorCopyTemplate(PyModuleState* modState, PyObject* tensor, const char* name = NULL, const char* dtype = NULL);
-PyObject* tensorCopyTemplateSimple(PyModuleState* modState, PyObject* tensor, const char* name_ = NULL, const char* dtype_ = NULL);
+PyObject* tensorCopyTemplate(PyModuleState* modState, PyObject* tensor, const char* name = NULL, PyObject* dtype = NULL);
+PyObject* tensorCopyTemplateSimple(
+    PyModuleState* modState, PyObject* tensor, const char* name_ = NULL, PyObject* dtype = NULL, bool copySparseDim = true);
 
 // exported Python functions {
 

{returnn-1.20231018.70522 → returnn-1.20231018.94805}/returnn/frontend/_utils.py

@@ -58,7 +58,7 @@ def bin_op_out_template(
     b: Union[Tensor[T], int, float, numpy.number],
     *,
     name: str,
-    res_dtype: Optional[str],
+    copy_sparse_dim: bool = True,
     allow_broadcast_all_sources: Optional[bool] = None,
     dim_order: Optional[Sequence[Dim]] = None,
     allow_scalar: bool = True,
@@ -70,7 +70,7 @@ def bin_op_out_template(
     :param a:
     :param b:
     :param name: for returned Tensor. no other functionality
-    :param res_dtype: if not given, infer from a and b
+    :param copy_sparse_dim:
     :param allow_broadcast_all_sources: if True, it is allowed that neither a nor b has all dims of the result.
         Not needed when out_dims is specified explicitly.
     :param dim_order: defines the order of the resulting dims. if None, it is automatically inferred from a and b.
@@ -93,10 +93,6 @@ def bin_op_out_template(
     # sanity checks
     # noinspection PyProtectedMember
     assert a._raw_backend == b._raw_backend, "Cannot combine tensors from two different frontends, e.g. TF and PT"
-    if res_dtype is None:
-        assert (
-            a.dtype == b.dtype
-        ), f"For now only operations with Tensors of the same dtypes are supported, got {a} and {b}"
     all_dims = []
     for dim in a.dims + b.dims:
         if dim in all_dims:
@@ -125,9 +121,9 @@ def bin_op_out_template(
             raise TypeError(f"invalid type for allow_broadcast_all_sources: {type(allow_broadcast_all_sources)}")
     if dim_order:
         all_dims.sort(key=lambda d: dim_order.index(d) if d in dim_order else len(dim_order))
-    out = Tensor(name, dims=all_dims, dtype=res_dtype or src_dtype)
+    out = Tensor(name, dims=all_dims, dtype=src_dtype)
     out.feature_dim = res_feature_dim(a, b)
-    if not res_dtype:
+    if copy_sparse_dim:
         out.sparse_dim = res_sparse_dim(a, b)
     if not allow_scalar or a.dims:
         a_raw = a.copy_compatible_to_dims_raw(all_dims)

{returnn-1.20231018.70522 → returnn-1.20231018.94805}/returnn/frontend/loop.py

@@ -236,10 +236,6 @@ def _check_matching_loop_var_templates(loop_var_templates: S, loop_vars: S):
                 f"loop var {path} template {template} does not match var {x}, "
                 f"different dims (no matter the order) {template.dims} vs {x.dims}"
             )
-            assert template.dtype == x.dtype, (
-                f"loop var {path} template {template} does not match var {x}, "
-                f"different dtype {template.dtype} vs {x.dtype}"
-            )
             assert template.sparse_dim == x.sparse_dim, (
                 f"loop var {path} template {template} does not match var {x}, "
                 f"different sparse_dim {template.sparse_dim} vs {x.sparse_dim}"

{returnn-1.20231018.70522 → returnn-1.20231018.94805}/returnn/frontend/tensor_array.py

@@ -94,10 +94,6 @@ class TensorArray:
             f"TensorArray push_back: template {self.tensor_template} does not match tensor {tensor},"
             f" dims different, {self.tensor_template.dims} vs {tensor.dims}"
         )
-        assert tensor.dtype == self.tensor_template.dtype, (
-            f"TensorArray push_back: template {self.tensor_template} does not match tensor {tensor},"
-            f" dtype different, {self.tensor_template.dtype} vs {tensor.dtype}"
-        )
         assert tensor.sparse_dim == self.tensor_template.sparse_dim, (
             f"TensorArray push_back: template {self.tensor_template} does not match tensor {tensor},"
             f" sparse_dim different, {self.tensor_template.sparse_dim} vs {tensor.sparse_dim}"

{returnn-1.20231018.70522 → returnn-1.20231018.94805}/returnn/torch/engine.py

@@ -95,7 +95,7 @@ class Engine(EngineBase):
             self._torch_distributed_class = torch_distributed.get("class", None)
             self._torch_distributed_options = torch_distributed.get("options", None)
 
-        amp_options = self.config.typed_value("torch_amp")
+        amp_options = self.config.opt_typed_value("torch_amp")
         grad_scaler_opts = self.config.typed_value("grad_scaler", NotSpecified)
         if amp_options is not None:
             self._use_autocast = True
@@ -284,7 +284,7 @@ class Engine(EngineBase):
                 step_idx % self._accum_grad_multiple_step
             ) != (self._accum_grad_multiple_step - 1) else nullcontext():
                 if self._grad_scaler is not None:
-                    self._grad_scaler.scale(total_loss).backward()
+                    self._grad_scaler.scale(total_loss.raw_tensor).backward()
                 else:
                     total_loss.raw_tensor.backward()
 

{returnn-1.20231018.70522 → returnn-1.20231018.94805}/returnn/torch/frontend/_backend.py

@@ -494,7 +494,9 @@ class TorchBackend(Backend[torch.Tensor]):
             mask = tensor.get_sequence_mask_broadcast(axis=axis)
             inf_value = get_global_inf_value()
             tensor.raw_tensor = torch.where(mask, tensor.raw_tensor, -inf_value)
-        out.raw_tensor = torch.softmax(tensor.raw_tensor, dim=tensor.dims.index(axis))
+        out_raw = torch.softmax(tensor.raw_tensor, dim=tensor.dims.index(axis))
+        out.dtype = TorchBackend.get_dtype_name_raw(out_raw)
+        out.raw_tensor = out_raw
         return out
 
     @staticmethod
@@ -511,7 +513,9 @@ class TorchBackend(Backend[torch.Tensor]):
             mask = tensor.get_sequence_mask_broadcast(axis=axis)
             inf_value = get_global_inf_value()
             tensor.raw_tensor = torch.where(mask, tensor.raw_tensor, -inf_value)
-        out.raw_tensor = torch.log_softmax(tensor.raw_tensor, dim=tensor.dims.index(axis))
+        out_raw = torch.log_softmax(tensor.raw_tensor, dim=tensor.dims.index(axis))
+        out.dtype = TorchBackend.get_dtype_name_raw(out_raw)
+        out.raw_tensor = out_raw
         return out
 
     @staticmethod
@@ -955,7 +959,6 @@ class TorchBackend(Backend[torch.Tensor]):
 
         if use_mask and any(dim.dyn_size_ext for dim in reduce):
             raise NotImplementedError("masking in matmul reduce not yet implemented")
-        assert a.dtype == b.dtype, f"matmul: dtypes do not match: {a} vs {b}"
 
         a_dims = a.dims
         b_dims = b.dims
@@ -1041,7 +1044,9 @@ class TorchBackend(Backend[torch.Tensor]):
         b_unique_dims = [b_dims[i] for i in b_unique_axes]
         result_dims = common_dims + a_unique_dims + b_unique_dims
 
-        result_tensor = Tensor(name="dot", dims=result_dims, raw_tensor=raw_result, dtype=a.dtype)
+        result_tensor = Tensor(
+            name="dot", dims=result_dims, raw_tensor=raw_result, dtype=TorchBackend.get_dtype_name_raw(raw_result)
+        )
 
         return result_tensor
 
@@ -1561,7 +1566,9 @@ class TorchBackend(Backend[torch.Tensor]):
             )
         else:
             raise ValueError(f"invalid number of filter dims {filter_size}, expected 1, 2, or 3")
-        out = Tensor("conv", dims=batch_dims + [out_dim] + list(out_spatial_dims), dtype=source.dtype)
+        out = Tensor(
+            "conv", dims=batch_dims + [out_dim] + list(out_spatial_dims), dtype=TorchBackend.get_dtype_name_raw(out_raw)
+        )
         if len(batch_dims) == 1:
             out.raw_tensor = out_raw
         else:

{returnn-1.20231018.70522 → returnn-1.20231018.94805/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20231018.70522
+Version: 1.20231018.94805
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer