returnn 1.20251027.232712__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()
@@ -206,6 +221,14 @@ Ndarray* Ndarray_Copy(const Ndarray* self) {
 
 #include "tensorflow/core/public/version.h"
 
+#ifndef TF_MAJOR_VERSION
+#error "TF_MAJOR_VERSION is not defined!"
+#endif
+
+#ifndef TF_MINOR_VERSION
+#error "TF_MINOR_VERSION is not defined!"
+#endif
+
 #if (TF_MAJOR_VERSION == 1 && TF_MINOR_VERSION >= 6) || (TF_MAJOR_VERSION > 1)
 #define TF_issue_6602_workaround 0
 #define TWOD_LSTM_SUPPORT 1
@@ -391,100 +414,172 @@ 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);
 
 
 #else  // CUDA
+
+#ifdef HAVE_CUSTOM_BLAS
+
 /*
     // 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
+
+template<typename T>
+static void tf_cpu_sgemm(
+    OpKernelContext* context,
+    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)
+{
+    if (m <= 0 || n <= 0 || k <= 0) return;
+
+    auto d = context->eigen_cpu_device();
+    const T alpha = *alpha_ptr;
+    const T beta = *beta_ptr;
+
+    bool transa = (transa_ == 'T' || transa_ == 't' || transa_ == 'C' || transa_ == 'c');
+    bool transb = (transb_ == 'T' || transb_ == 't' || transb_ == 'C' || transb_ == 'c');
+
+    // 1. Map as COLUMN-MAJOR
+    // Physical rows (height) for the Map is always the leading dimension (lda, ldb, ldc)
+    typedef Eigen::TensorMap<Eigen::Tensor<const T, 2, Eigen::ColMajor>, Eigen::Unaligned> ConstMap;
+    typedef Eigen::TensorMap<Eigen::Tensor<T, 2, Eigen::ColMajor>, Eigen::Unaligned> MutableMap;
+
+    // Logical height/width of slices before any transposition
+    int a_slice_rows = transa ? k : m;
+    int a_slice_cols = transa ? m : k;
+    int b_slice_rows = transb ? n : k;
+    int b_slice_cols = transb ? k : n;
+
+    // Map and Slice
+    auto a = ConstMap(a_ptr, lda, a_slice_cols).slice(
+        Eigen::array<Eigen::Index, 2>({0, 0}),
+        Eigen::array<Eigen::Index, 2>({(Eigen::Index)a_slice_rows, (Eigen::Index)a_slice_cols}));
+
+    auto b = ConstMap(b_ptr, ldb, b_slice_cols).slice(
+        Eigen::array<Eigen::Index, 2>({0, 0}),
+        Eigen::array<Eigen::Index, 2>({(Eigen::Index)b_slice_rows, (Eigen::Index)b_slice_cols}));
+
+    auto c = MutableMap(c_ptr, ldc, n).slice(
+        Eigen::array<Eigen::Index, 2>({0, 0}),
+        Eigen::array<Eigen::Index, 2>({(Eigen::Index)m, (Eigen::Index)n}));
+
+    // 2. Define Contraction Pairs based on Transposition
+    // Column-Major Matrix Mult: (M x K) * (K x N)
+    // Standard: Contract Axis 1 of A with Axis 0 of B
+    // If A is Transposed: A is (K x M), contract Axis 0 of A
+    // If B is Transposed: B is (N x K), contract Axis 1 of B
+    Eigen::array<Eigen::IndexPair<int>, 1> pairs;
+    pairs[0] = Eigen::IndexPair<int>(transa ? 0 : 1, transb ? 1 : 0);
+
+    // 3. Execution
+    if (alpha == T(1) && beta == T(0)) {
+        c.device(d) = a.contract(b, pairs);
+    } else if (alpha == T(1) && beta == T(1)) {
+        c.device(d) += a.contract(b, pairs);
+    } else {
+        c.device(d) = a.contract(b, pairs) * alpha + c * beta;
+    }
+}
+
+#define Ndarray_sgemm(\
+    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"
 
-#if CUDA
+#endif // TENSORFLOW or TORCH
 
 
+#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))
@@ -501,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__ ))
@@ -550,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))
@@ -671,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;
 }
@@ -769,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/sprint/cache.py

@@ -7,10 +7,9 @@ This module is about reading (maybe later also writing) the Sprint archive forma
 """
 
 from __future__ import annotations
-from typing import List, Optional, Tuple
+from typing import Optional, List, Tuple, Dict
 import sys
 import os
-import typing
 import array
 from struct import pack, unpack
 import numpy
@@ -212,7 +211,7 @@ class FileArchive:
     def __init__(self, filename, must_exists=True, encoding="ascii"):
         self.encoding = encoding
 
-        self.ft = {}  # type: typing.Dict[str,FileInfo]
+        self.ft: Dict[str, FileInfo] = {}
         if os.path.exists(filename):
             self.allophones = []
             self.f = open(filename, "rb")
@@ -334,8 +333,8 @@ class FileArchive:
             # print(typ)
             assert type_ == "vector-f32"
             count = self.read_U32()
-            data = [None] * count  # type: typing.List[typing.Optional[numpy.ndarray]]
-            time_ = [None] * count  # type: typing.List[typing.Optional[numpy.ndarray]]
+            data: List[Optional[numpy.ndarray]] = [None] * count
+            time_: List[Optional[numpy.ndarray]] = [None] * count
             for i in range(count):
                 size = self.read_U32()
                 data[i] = self.read_v("f", size)  # size x f32
@@ -450,7 +449,7 @@ class FileArchive:
             a = array.array("b")
             a.fromfile(self.f, comp)
             # unpack
-            b = zlib.decompress(a.tostring(), 15 + 32)
+            b = zlib.decompress(a.tobytes(), 15 + 32)
             # substitute self.f by an anonymous memmap file object
             # restore original file handle after we're done
             backup_f = self.f
@@ -575,17 +574,17 @@ class FileArchiveBundle:
         :param str encoding: encoding used in the files
         """
         # filename -> FileArchive
-        self.archives = {}  # type: typing.Dict[str,FileArchive]
+        self.archives: Dict[str, FileArchive] = {}
         # archive content file -> FileArchive
-        self.files = {}  # type: typing.Dict[str,FileArchive]
+        self.files: Dict[str, FileArchive] = {}
         self._short_seg_names = {}
         if filename is not None:
             self.add_bundle(filename=filename, encoding=encoding)
 
-    def add_bundle(self, filename, encoding="ascii"):
+    def add_bundle(self, filename: str, encoding: str = "ascii"):
         """
-        :param str filename: bundle
-        :param str encoding:
+        :param filename: bundle
+        :param encoding:
         """
         file_dir = os.path.dirname(filename) or "."
         for line in open(filename).read().splitlines():
@@ -837,7 +836,7 @@ class MixtureSet:
         """
         a = array.array("b")
         a.fromfile(self.f, length)
-        return a.tostring().decode(encoding)
+        return a.tobytes().decode(encoding)
 
     def read_f32(self):
         """
@@ -1003,7 +1002,7 @@ class WordBoundaries:
         """
         a = array.array("b")
         a.fromfile(self.f, length)
-        return a.tostring().decode(encoding)
+        return a.tobytes().decode(encoding)
 
     def __init__(self, filename):
         """