returnn 1.20260105.192646__py3-none-any.whl → 1.20260119.15400__py3-none-any.whl

returnn/native_op.cpp

@@ -1,3 +1,14 @@
+/*
+This file is imported in various ways.
+The mode is determined via the preprocessor defines:
+
+TENSORFLOW: If defined and set to 1, TensorFlow is used as backend.
+TORCH: If defined and set to 1, PyTorch is used as backend.
+
+CUDA: If defined and set to 1, CUDA is used for GPU support.
+    Otherwise, it uses CPU only.
+    The kernels are all expected to also compile in CPU-only mode.
+*/
 
 #include <assert.h>
 #include <iostream>
@@ -16,6 +27,10 @@
 #define TENSORFLOW 0
 #endif
 
+#ifndef TORCH
+#define TORCH 0
+#endif
+
 #ifndef _ns
 #define _ns
 #endif
@@ -118,7 +133,7 @@ static inline int _host_float_as_int(float x) {
 #define INF_F int_as_float(0x7f800000)
 #define NAN_F int_as_float(0x7fffffff)
 
-#endif
+#endif // CUDA
 
 
 
@@ -157,7 +172,7 @@ The BLAS functions expect the inputs in column-major and return in column-major.
 #define Ndarray tensorflow::Tensor
 #define Ndarray_DEV_DATA(x) ((float*) (x)->tensor_data().data())
 #define Ndarray_DEV_DATA_int32(x) ((int32_t*) (x)->tensor_data().data())
-#define Ndarray_DEV_DATA_int32_scalar(x) (x)->scalar<int32>()()
+#define Ndarray_DEV_DATA_int32_scalar(x) (x)->scalar<int32_t>()()
 #define Ndarray_HOST_DIMS(x) DimsAccessor(x)
 #define Ndarray_DIMS Ndarray_HOST_DIMS
 #define Ndarray_NDIM(x) (x)->dims()
@@ -399,13 +414,13 @@ static void tf_cuda_sgemm_batched(
 
 
 #define Ndarray_sgemm( \
-	transpose_A, transpose_B, \
-	m, n, k, alpha, A, lda, B, ldb, beta, C, ldc) \
+    transpose_A, transpose_B, \
+    m, n, k, alpha, A, lda, B, ldb, beta, C, ldc) \
     tf_cuda_sgemm<float>(context, transpose_A, transpose_B, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
 
 #define Ndarray_sgemm_batched( \
-	transpose_A, transpose_B, \
-	m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, batchSize, finalize_stream) \
+    transpose_A, transpose_B, \
+    m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, batchSize, finalize_stream) \
     tf_cuda_sgemm_batched<float>(context, transpose_A, transpose_B, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, batchSize, finalize_stream);
 
 
@@ -415,21 +430,21 @@ static void tf_cuda_sgemm_batched(
 
 /*
     // matrices are in column-major form
-	int sgemm_(char *transa, char *transb,
-		integer *m, integer *n, integer *k,
-		real *alpha, real *a, integer *lda,
-		real *b, integer *ldb, real *beta,
-		real *c, integer *ldc);
+    int sgemm_(char *transa, char *transb,
+        integer *m, integer *n, integer *k,
+        real *alpha, real *a, integer *lda,
+        real *b, integer *ldb, real *beta,
+        real *c, integer *ldc);
 */
 #define Ndarray_sgemm(\
-	transpose_A, transpose_B, \
-	m, n, k, alpha, A, lda, B, ldb, beta, C, ldc) \
-	{ \
-		char transa = transpose_A, transb = transpose_B; \
-		int m_ = m, n_ = n, k_ = k, lda_ = lda, ldb_ = ldb, ldc_ = ldc; \
-		sgemm_(&transa, &transb, \
-			&m_, &n_, &k_, alpha, A, &lda_, B, &ldb_, beta, C, &ldc_); \
-	}
+    transpose_A, transpose_B, \
+    m, n, k, alpha, A, lda, B, ldb, beta, C, ldc) \
+    { \
+        char transa = transpose_A, transb = transpose_B; \
+        int m_ = m, n_ = n, k_ = k, lda_ = lda, ldb_ = ldb, ldc_ = ldc; \
+        sgemm_(&transa, &transb, \
+            &m_, &n_, &k_, alpha, A, &lda_, B, &ldb_, beta, C, &ldc_); \
+    }
 
 #else  // HAVE_CUSTOM_BLAS
 
@@ -494,77 +509,77 @@ static void tf_cpu_sgemm(
 }
 
 #define Ndarray_sgemm(\
-	transpose_A, transpose_B, \
-	m, n, k, alpha, A, lda, B, ldb, beta, C, ldc) \
+    transpose_A, transpose_B, \
+    m, n, k, alpha, A, lda, B, ldb, beta, C, ldc) \
     tf_cpu_sgemm<float>(context, transpose_A, transpose_B, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
 
 #endif  // HAVE_CUSTOM_BLAS
 #endif  // CUDA
 
+#define CHECK_WITH_MSG(condition, message) \
+    if(!(condition)) { \
+        std::cerr << "NativeOp check failed: " << message << std::endl; \
+        assert(condition); \
+    }
+
 // See Context struct below.
 #define CONTEXT_ARGS    context
 
-#else  // TENSORFLOW
+
+#elif TORCH
+// https://github.com/rwth-i6/i6_native_ops/blob/main/i6_native_ops/common/returnn_definitions.h
+// https://docs.pytorch.org/cppdocs/stable.html#tensor-class
+
+#define Ndarray torch::Tensor
+#define Ndarray_DEV_DATA(x) ((float*)(x)->data_ptr())
+#define Ndarray_DEV_DATA_int32(x) ((int32_t*)(x)->data_ptr())
+#define Ndarray_DEV_DATA_uint32(x) ((uint32_t*)(x)->data_ptr())
+#define Ndarray_DEV_DATA_int32_scalar(x) ((x)->item().to<int32_t>())
+#define Ndarray_HOST_DIMS(x) ((x)->sizes())
+#define Ndarray_DIMS(x) ((x)->sizes())
+typedef at::IntArrayRef Ndarray_DIMS_Type;
+#define Ndarray_NDIM(x) (x)->dim()
+#define Ndarray_dtype_size(x) torch::elementSize((x)->scalar_type())
+typedef int64_t Ndarray_DIM_Type;
+#define Ndarray_SIZE(x) ((x)->numel())
+#define Ndarray_STRIDE(x, dim) ((x)->stride(dim))
+
+#define CHECK_WITH_MSG TORCH_CHECK
 
 // See Context struct below.
 #define CONTEXT_ARGS
 
-#endif  // TENSORFLOW
+template<typename T>
+static void Ndarray_sgemm(
+    char transa_, char transb_,
+    int m, int n, int k,
+    const T* alpha_ptr, const T* a_ptr, int lda,
+    const T* b_ptr, int ldb, const T* beta_ptr,
+    T* c_ptr, int ldc)
+{
+    // TODO...
+    assert("Torch Ndarray_sgemm not implemented" && 0);
+}
 
+#else  // TENSORFLOW or TORCH
 
+#error "No framework defined: TENSORFLOW or TORCH"
+
+#endif // TENSORFLOW or TORCH
 
-#if CUDA
 
+#if CUDA
 
 #if TENSORFLOW
 // Ndarray and friends already declared above, they are same for CUDA and non-CUDA
 #define CUDA_CUR_STREAM  (context->eigen_gpu_device().stream())
 
-#else  // TENSORFLOW, thus Theano here
-#define CUDA_CUR_STREAM  (0)  // default stream
+#elif TORCH
 
-// Defined here: https://github.com/Theano/Theano/blob/master/theano/sandbox/cuda/cuda_ndarray.cuh
-// See also: https://github.com/Theano/Theano/blob/master/theano/sandbox/cuda/cuda_ndarray.cu
-#define Ndarray CudaNdarray
-#define Ndarray_DEV_DATA CudaNdarray_DEV_DATA
-#define Ndarray_DEV_DATA_int32(x) ((int32_t*) (Ndarray_DEV_DATA(x)))
-#define Ndarray_DEV_DATA_int32_scalar(x) Ndarray_DEV_DATA_int32(x)[0]
-#define Ndarray_HOST_DIMS CudaNdarray_HOST_DIMS
-#define Ndarray_DIMS Ndarray_HOST_DIMS
-#define Ndarray_STRIDE(x, i) (CudaNdarray_HOST_STRIDES(x)[i])  // return in elements. CudaNdarray stores like that
-#define Ndarray_NDIM(x) (x->nd)
-#define Ndarray_DIM_Type int
-typedef Ndarray_DIM_Type const* Ndarray_DIMS_Type;
-#define Ndarray_dtype_size(x) sizeof(float)
-#define Ndarray_SIZE CudaNdarray_SIZE
-// PyObject *CudaNdarray_NewDims(int nd, const inttype * dims), uninitialized
-#define Ndarray_NewDims CudaNdarray_NewDims
-// PyObject * CudaNdarray_Copy(const CudaNdarray * self);
-#define Ndarray_Copy CudaNdarray_Copy
-
-/*
-    // via: https://docs.nvidia.com/cuda/cublas/
-    // matrices are in column-major form
-    cublasStatus_t cublasSgemm(cublasHandle_t handle,
-        cublasOperation_t transa, cublasOperation_t transb,
-        int m, int n, int k,
-        const float *alpha, const float *A, int lda,
-        const float *B, int ldb, const float *beta,
-        float *C, int ldc);
-*/
-#define _cublasTranspose(t) \
-	((t == 'T') ? CUBLAS_OP_T : \
-	(t == 'C') ? CUBLAS_OP_C : \
-	(t == 'N') ? CUBLAS_OP_N : cublasOperation_t('E'))
-#define Ndarray_sgemm( \
-	transpose_A, transpose_B, \
-	m, n, k, alpha, A, lda, B, ldb, beta, C, ldc) \
-	(_cudaHandleError(cublasSgemm(handle, \
-	_cublasTranspose(transpose_A), \
-	_cublasTranspose(transpose_B), \
-	m, n, k, alpha, A, lda, B, ldb, beta, C, ldc), \
-	__FILE__, __LINE__ ))
+#define CUDA_CUR_STREAM (at::cuda::getCurrentCUDAStream().stream())
 
+#else
+#error Unknown backend
 #endif
 
 #define Ndarray_memcpy(y, x, size) (cudaMemcpyAsync(y, x, size, cudaMemcpyDeviceToDevice, CUDA_CUR_STREAM))
@@ -581,48 +596,10 @@ typedef Ndarray_DIM_Type const* Ndarray_DIMS_Type;
 
 #define DEF_SHARED(type, name) extern __shared__ type name[];
 
-static const char *_cudaGetErrorEnum(cublasStatus_t error) {
-	switch (error) {
-	case CUBLAS_STATUS_SUCCESS:
-		return "CUBLAS_STATUS_SUCCESS";
-
-	case CUBLAS_STATUS_NOT_INITIALIZED:
-		return "CUBLAS_STATUS_NOT_INITIALIZED";
-
-	case CUBLAS_STATUS_ALLOC_FAILED:
-		return "CUBLAS_STATUS_ALLOC_FAILED";
-
-	case CUBLAS_STATUS_INVALID_VALUE:
-		return "CUBLAS_STATUS_INVALID_VALUE";
-
-	case CUBLAS_STATUS_ARCH_MISMATCH:
-		return "CUBLAS_STATUS_ARCH_MISMATCH";
-
-	case CUBLAS_STATUS_MAPPING_ERROR:
-		return "CUBLAS_STATUS_MAPPING_ERROR";
-
-	case CUBLAS_STATUS_EXECUTION_FAILED:
-		return "CUBLAS_STATUS_EXECUTION_FAILED";
-
-	case CUBLAS_STATUS_INTERNAL_ERROR:
-		return "CUBLAS_STATUS_INTERNAL_ERROR";
-	}
-
-	return "<unknown>";
-}
-
-static void _cudaHandleError(cudaError_t err, const char *file, int line) {
-	if (err != cudaSuccess) {
-		printf("NativeOp: CUDA runtime error: '%s' in %s at line %d\n", cudaGetErrorString(err), file, line);
-		exit(EXIT_FAILURE);
-	}
-}
-
-static void _cudaHandleError(cublasStatus_t status, const char *file, int line) {
-	if (status != CUBLAS_STATUS_SUCCESS) {
-		printf("NativeOp: cuBLAS runtime error: '%s' in %s at line %d\n", _cudaGetErrorEnum(status), file, line);
-		exit(EXIT_FAILURE);
-	}
+static void _cudaHandleError(cudaError_t err, const char* file, int line) {
+    CHECK_WITH_MSG(
+        err == cudaSuccess,
+        "NativeOp: CUDA runtime error: ", cudaGetErrorString(err), " in ", file, " at line ", line);
 }
 
 #define HANDLE_ERROR(status) (_cudaHandleError( status, __FILE__, __LINE__ ))
@@ -630,49 +607,7 @@ static void _cudaHandleError(cublasStatus_t status, const char *file, int line)
 
 #else   // not CUDA
 
-
-#if !TENSORFLOW
-// Numpy, see: https://docs.scipy.org/doc/numpy/reference/c-api.array.html
-// And: https://deeplearning.net/software/theano/extending/extending_theano_c.html
-#define Ndarray PyArrayObject
-#define Ndarray_DEV_DATA(x) ((float*) PyArray_DATA(x))
-#define Ndarray_DEV_DATA_int32(x) ((int32_t*) (Ndarray_DEV_DATA(x)))
-#define Ndarray_DEV_DATA_int32_scalar(x) Ndarray_DEV_DATA_int32(x)[0]
-#define Ndarray_HOST_DIMS PyArray_DIMS
-#define Ndarray_STRIDE(x, i) (PyArray_STRIDE(x, i) / sizeof(float))  // return in elements. Numpy stores in bytes
-#define Ndarray_DIMS Ndarray_HOST_DIMS
-#define Ndarray_NDIM PyArray_NDIM
-#define Ndarray_DIM_Type npy_intp
-typedef Ndarray_DIM_Type const* Ndarray_DIMS_Type;
-#define Ndarray_dtype_size(x) sizeof(float)
-#define Ndarray_SIZE PyArray_SIZE
-#define Ndarray_NewDims(nd, dims) (PyArray_SimpleNew(nd, dims, NPY_FLOAT32))
-#define Ndarray_Copy(x) (PyArray_FromArray(x, NULL, NPY_ARRAY_OUT_ARRAY | NPY_ARRAY_ENSURECOPY))
-/*
-    // matrices are in column-major form
-	int sgemm_(char *transa, char *transb,
-		integer *m, integer *n, integer *k,
-		real *alpha, real *a, integer *lda,
-		real *b, integer *ldb, real *beta,
-		real *c, integer *ldc);
-
-	Cast to (float*) because we might have the C-style declaration incorrectly in the C++ scope.
-*/
-#define Ndarray_sgemm(\
-	transpose_A, transpose_B, \
-	m, n, k, alpha, A, lda, B, ldb, beta, C, ldc) \
-	{ \
-		char transa = transpose_A, transb = transpose_B; \
-		int m_ = m, n_ = n, k_ = k, lda_ = lda, ldb_ = ldb, ldc_ = ldc; \
-		sgemm_(&transa, &transb, \
-			&m_, &n_, &k_, alpha, (float*) A, &lda_, (float*) B, &ldb_, beta, C, &ldc_); \
-	}
-
-static inline void* device_malloc(size_t size) { return malloc(size); }
-static inline void device_free(void* ptr) { free(ptr); }
-#endif
-
-#define HANDLE_LAST_ERROR() (0)
+#define HANDLE_LAST_ERROR() {}
 
 #define Ndarray_memcpy(y, x, size) (memcpy(y, x, size))
 #define Ndarray_memset(s, c, size) (memset(s, c, size))
@@ -751,19 +686,9 @@ struct _KernelLoop {
 #endif
 
 
-Ndarray* Ndarray_uninitialized_like(Ndarray* a) {
-	Ndarray_DIMS_Type dim = Ndarray_HOST_DIMS(a);
-#if TENSORFLOW
-	Ndarray* res = (Ndarray*) Ndarray_NewDims(Ndarray_NDIM(a), dim);
-#else
-	Ndarray* res = (Ndarray*) Ndarray_NewDims(Ndarray_NDIM(a), const_cast<Ndarray_DIM_Type*>(dim));
-#endif
-	return res;
-}
-
-long Ndarray_get_n_total_elements(Ndarray* a) {
-	long c = 1;
-	for(long i = 0; i < Ndarray_NDIM(a); ++i)
+int64_t Ndarray_get_n_total_elements(Ndarray* a) {
+	int64_t c = 1;
+	for(int i = 0; i < Ndarray_NDIM(a); ++i)
 		c *= Ndarray_DIMS(a)[i];
 	return c;
 }
@@ -849,17 +774,22 @@ void _free(void* ptr) {
         context->device()->GetAllocator(AllocatorAttributes());
     allocator->DeallocateRaw(ptr);
 }
-#define device_malloc Context(CONTEXT_ARGS)._malloc
-#define device_free Context(CONTEXT_ARGS)._free
+
+#elif TORCH
 
 #if CUDA
-cublasHandle_t _handle() {
-    assert("not available" && 0);
-    return NULL;
-}
-#define handle Context(CONTEXT_ARGS)._handle()
-#endif
-#endif
+void* _malloc(size_t num_bytes) { return c10::cuda::CUDACachingAllocator::raw_alloc(num_bytes); }
+void _free(void* ptr) { c10::cuda::CUDACachingAllocator::raw_delete(ptr); }
+#else  // not CUDA
+void* _malloc(size_t num_bytes) { return c10::GetCPUAllocator()->raw_allocate(num_bytes); }
+void _free(void* ptr) { c10::GetCPUAllocator()->raw_deallocate(ptr); }
+#endif  // CUDA
+
+#endif  // TENSORFLOW or TORCH
+
+
+#define device_malloc Context(CONTEXT_ARGS)._malloc
+#define device_free Context(CONTEXT_ARGS)._free
 
 
 //C[x] += A[x]*B[x]

returnn/native_op.py

@@ -5,38 +5,40 @@ Generic interface which automatically creates:
 * inplace and not inplace
 * grad variants
 
-See :mod:`returnn.tf.native_op` and :mod:`returnn.theano.native_op`
-for usage in TensorFlow and Theano.
+See :mod:`returnn.tf.native_op` and :mod:`returnn.torch.utils.native_op`
+for usage in TensorFlow and PyTorch.
 
 See :ref:`native_ops` for more background.
 """
 
+from __future__ import annotations
+from typing import Optional, Union, Any, Callable, Dict, Sequence, Tuple
 import copy
 import numpy
-import typing
+
 from returnn.util.basic import make_hashable, unicode
 
 
 class NativeOpBaseMixin:
     """
-    The purpose of having this as a separate base class is to make this independent of any Theano specific
-    functionality so that we can also use this base for example for TensorFlow.
+    The purpose of having this as a separate base class
+    is to make this independent of any TensorFlow or PyTorch-specific functionality.
     """
 
     def __init__(
         self,
-        in_info,
-        out_info,
-        c_fw_code,
-        c_bw_code=None,
-        c_extra_support_code=None,
-        code_version=None,
-        cpu_support=True,
-        grad_input_map=None,
-        name=None,
+        in_info: Sequence[Dict[str, Any]],
+        out_info: Sequence[Dict[str, Any]],
+        c_fw_code: str,
+        c_bw_code: Optional[str] = None,
+        c_extra_support_code: Union[None, str, Dict[str, str]] = None,
+        code_version: Optional[Tuple[int, ...]] = None,
+        cpu_support: bool = True,
+        grad_input_map: Union[None, Tuple[int, ...], Callable] = None,
+        name: Optional[str] = None,
     ):
         """
-        :param list[dict(str)] in_info: each dict describes one input var.
+        :param in_info: each dict describes one input var.
           attribs in the dict:
             int ndim: the ndim.
             tuple shape: tuple and can contain None for specific dimensions.
@@ -49,18 +51,18 @@ class NativeOpBaseMixin:
             str gradient: can be "disconnected". see grad().
             bool bw_input: True by default. add this param to the bw input.
           other attribs are just ignored.
-        :param list[dict(str)] out_info: like in_info.
+        :param out_info: like in_info.
           slightly different behavior for:
             shape: we also allow refs to the in_info in the form (in-idx,dim). see infer_shape().
             need_contiguous/want_inplace: used for bw, in case for bw_input == True.
-        :param str c_fw_code: C code for forward pass
-        :param str|dict[str] c_extra_support_code: C support code (for c_support_code)
-        :param str|None c_bw_code: C code for backward pass (for gradient)
-        :param tuple[int] code_version: will be returned by c_code_cache_version.
-        :param bool cpu_support:
-        :param tuple[int]|callable grad_input_map: selection of grad inputs.
+        :param c_fw_code: C code for forward pass
+        :param c_extra_support_code: C support code (for c_support_code)
+        :param c_bw_code: C code for backward pass (for gradient)
+        :param code_version: will be returned by c_code_cache_version.
+        :param cpu_support:
+        :param grad_input_map: selection of grad inputs.
           by default, we get all inputs + all outputs + all grad outputs.
-        :param str name: name
+        :param name: name
         """
         assert isinstance(in_info, (list, tuple))
         assert isinstance(out_info, (list, tuple))
@@ -251,12 +253,12 @@ class NativeOpGenBase:
     See NativeOp.__init__() for attribs.
     """
 
-    in_info = None  # type: typing.Tuple[typing.Dict[str]]
-    out_info = None  # type: typing.Tuple[typing.Dict[str]]
-    c_fw_code = None  # type: str
-    c_bw_code = None  # type: str
-    c_extra_support_code = None  # type: typing.Dict[str,str]
-    code_version = None  # type: typing.Union[typing.Tuple[int], int]
+    in_info: Optional[Tuple[Dict[str, Any], ...]] = None
+    out_info: Optional[Tuple[Dict[str, Any], ...]] = None
+    c_fw_code: Optional[str] = None
+    c_bw_code: Optional[str] = None
+    c_extra_support_code: Optional[Dict[str, str]] = None
+    code_version: Union[None, Tuple[int, ...], int] = None
     grad_input_map = None
     theano_custom_grad = None
     cpu_support = True
@@ -4699,7 +4701,7 @@ class FastViterbiOp(NativeOpGenBase):
         int n_states,
         int n_edges,
         int t,
-        int32* cur_state, // (n_batch,)
+        int32_t* cur_state, // (n_batch,)
         const IdxAndVal* frame,
         const int32_t* d_am_seq_len,
         const int32_t* d_edge_from,
@@ -5339,7 +5341,10 @@ class EditDistanceOp(NativeOpGenBase):
             sub_cost = last1_dist[last1_idx];
             if(a[batch_idx * n_a_max_len + t_a - 1] != b[batch_idx * n_b_max_len + t_b - 1])
               ++sub_cost;
-            //printf("t_a %i, t_b %i, del %i, ins %i, sub %i\\n", t_a, t_b, del_cost, ins_cost, sub_cost);
+            /*printf("t_a %i, t_b %i, a %d, b %d, del %i, ins %i, sub %i\\n",
+                t_a, t_b,
+                a[batch_idx * n_a_max_len + t_a - 1], b[batch_idx * n_b_max_len + t_b - 1],
+                del_cost, ins_cost, sub_cost);*/
             int min_cost = del_cost;
             if(min_cost > ins_cost) min_cost = ins_cost;
             if(min_cost > sub_cost) min_cost = sub_cost;

returnn/tensor/_dim_extra.py

@@ -858,7 +858,7 @@ class _DimMixin:
         self._make_extra()
         dim_order_default = self.dyn_size_ext.dims + (self,)
         if dim_order is not None:
-            dim_order = tuple(d for d in dim_order if d in dim_order_default)  # filter
+            dim_order = tuple([d for d in dim_order if d in dim_order_default])  # filter
         else:
             dim_order = dim_order_default
         cache_key = (device, dim_order)
@@ -2484,16 +2484,16 @@ _BinOpStrs = {
 
 def _math_get_dim_via_bin_op(dims: Sequence[Union[Dim, int]], op_kind: str) -> Dim:
     dims = [d if isinstance(d, _d.Dim) else _make_constant_static_dim(d) for d in dims]
-    if all(d.dimension is not None for d in dims):
+    if all([d.dimension is not None for d in dims]):
         op = _BinOps[op_kind]
         dim_value = dims[0].dimension
         for d in dims[1:]:
             dim_value = op(dim_value, d.dimension)
     else:
         dim_value = None
-    if all(d.is_constant_static_dim() for d in dims):
+    if all([d.is_constant_static_dim() for d in dims]):
         return _make_constant_static_dim(dim_value, kind=_get_merged_dim_kind(dims))
-    desc = _BinOpStrs[op_kind].join(_get_description(d) for d in dims)
+    desc = _BinOpStrs[op_kind].join([_get_description(d) for d in dims])
     if op_kind.startswith("ceildiv"):
         desc = f"⌈{desc}⌉"
     return _d.Dim(
@@ -2676,16 +2676,16 @@ def _get_description(dim, brackets=True):
 
 
 def _get_merged_dim_kind(dim_tags: Sequence[Dim]) -> Entity:
-    if any(tag.is_batch_dim() for tag in dim_tags):
+    if any([tag.is_batch_dim() for tag in dim_tags]):
         return DimTypes.Batch
-    elif any(tag.is_feature_dim() for tag in dim_tags):
+    elif any([tag.is_feature_dim() for tag in dim_tags]):
         return DimTypes.Feature
     else:
         return DimTypes.Spatial
 
 
 def _representative_tag(terms: Sequence[Dim]) -> Optional[Dim]:
-    if any(not term_.auto_generated for term_ in terms):
+    if any([not term_.auto_generated for term_ in terms]):
         # Always prefer non-auto-generated.
         terms = [term_ for term_ in terms if not term_.auto_generated]
     # First find any dynamic.

returnn/tensor/_tensor_extra.py

@@ -32,8 +32,8 @@ class _TensorExtra:
         tensor: Tensor,
         time_dim_axis=NotSpecified,
         available_for_inference=True,
-        batch=None,
-        beam=None,
+        batch: Optional[BatchInfo] = None,
+        beam: Optional[SearchBeam] = None,
         control_flow_ctx=None,
     ):
         """
@@ -41,8 +41,8 @@ class _TensorExtra:
         :param int|None|NotSpecified time_dim_axis: where we have the time dim axis, after we added the batch-dim.
             this is often 1. however, can be None if there is no time-dim.
         :param bool available_for_inference: e.g. the extern data "classes" is usually not available for inference
-        :param BatchInfo|None batch:
-        :param SearchBeam|None beam: the batch-dim could be extended by a beam-size,
+        :param batch:
+        :param beam: the batch-dim could be extended by a beam-size,
             such that it represents the merged dims [batch, beam_size].
         :param ControlFlowContext|None control_flow_ctx:
         """
@@ -668,11 +668,11 @@ class _TensorMixin(_TensorMixinBase):
         if not perm:
             return self.copy()
         if allow_int and isinstance(perm[0], int):
-            assert all(isinstance(a, int) for a in perm), f"{self}: invalid perm {perm!r} types"
+            assert all([isinstance(a, int) for a in perm]), f"{self}: invalid perm {perm!r} types"
             assert set(perm) == set(range(len(perm))), f"{self}: invalid perm {perm!r}"
             return self._copy_compatible_to_dims_with_perm([self._dims[i] for i in perm], perm)
         else:
-            assert all(isinstance(a, Dim) for a in perm), f"{self}: invalid perm {perm!r} types"
+            assert all([isinstance(a, Dim) for a in perm]), f"{self}: invalid perm {perm!r} types"
             return self.copy_compatible_to_dims(perm)
 
     def copy_move_axis(self, old_axis, new_axis) -> _t.Tensor:
@@ -1155,7 +1155,7 @@ class _TensorMixin(_TensorMixinBase):
                     )
 
         assert v.batch_ndim == data.batch_ndim
-        assert all(mapped_axes[ax] == ax for ax in range(v.batch_ndim))
+        assert all([mapped_axes[ax] == ax for ax in range(v.batch_ndim)])
 
         if self.version == 1:
             # Ensure time_dim_axis and feature_dim_axis is same as in data
@@ -1702,7 +1702,7 @@ class _TensorMixin(_TensorMixinBase):
         """
         :return: shape with added batch-dim. e.g. (batch,time,feat) = (None,None,128)
         """
-        return tuple(tag.dimension for tag in self.dim_tags)
+        return tuple([tag.dimension for tag in self.dim_tags])
 
     # noinspection PyShadowingNames
     def get_batch_shape(self, batch_dim):
@@ -3214,7 +3214,7 @@ class _TensorMixin(_TensorMixinBase):
         if len(sources) == 1:
             return sources[0].copy_template()
         max_ndim = max([s.batch_ndim for s in sources])
-        if any(src.batch for src in sources):
+        if any([src.batch for src in sources]):
             from returnn.tf.util.data import BatchInfo
 
             common_batch = BatchInfo.get_common_batch_info([src.batch for src in sources if src.batch])
@@ -3254,7 +3254,7 @@ class _TensorMixin(_TensorMixinBase):
             else:
                 axis = common.get_default_new_axis_for_dim_tag(dim_tag)
                 common = common.copy_add_dim_by_tag(dim_tag, unbroadcast=True, axis=axis)
-        if all(s.batch_ndim < common.batch_ndim for s in sources):
+        if all([s.batch_ndim < common.batch_ndim for s in sources]):
             from returnn.util.basic import validate_broadcast_all_sources
 
             validate_broadcast_all_sources(

returnn/tensor/utils.py

@@ -71,7 +71,7 @@ def tensor_fill_random_numpy_(
                     # Make sure at least one of the dyn sizes matches the max size.
                     i = rnd.randint(0, dim.dyn_size_ext.raw_tensor.size)
                     dim.dyn_size_ext.raw_tensor.flat[i] = dyn_dim_max_sizes[dim]
-                    if dim in dyn_dim_min_sizes:
+                    if dim in dyn_dim_min_sizes and dim.dyn_size_ext.raw_tensor.size > 1:
                         j = rnd.randint(0, dim.dyn_size_ext.raw_tensor.size - 1)
                         if j >= i:
                             j += 1

returnn/tf/frontend_layers/_backend.py

@@ -465,6 +465,8 @@ class ReturnnLayersBackend(Backend[Layer]):
         targets_spatial_dim: Dim,
         blank_index: int,
         max_approx: bool = False,
+        use_native_op: Optional[bool] = None,
+        label_loop: bool = True,
     ) -> Tensor:
         """CTC"""
         assert targets.sparse_dim and targets.sparse_dim.dimension <= logits.feature_dim.dimension
@@ -482,6 +484,7 @@ class ReturnnLayersBackend(Backend[Layer]):
                 "targets": targets,
                 "blank_index": blank_index,
                 "max_approx": max_approx,
+                "label_loop": label_loop,
             },
             name="ctc_loss",
         )
@@ -944,7 +947,6 @@ class ReturnnLayersBackend(Backend[Layer]):
         """
         assert mask.dtype == "bool"
         assert set(mask.dims) == set(dims)
-        assert set(mask.dims).issubset(set(tensor.dims))
         if not out_dim:
             out_dim = Dim(None, name="mask")
         return (