<commit_before>import numpy
import warnings

import cupy

from libc.stdint cimport intptr_t, uint32_t, uint64_t
from cupy.core._carray cimport shape_t
from cupy.core.core cimport ndarray
from cupy.core cimport internal
from cupy_backends.cuda.libs.cutensor cimport Handle
from cupy_backends.cuda.libs.cutensor cimport TensorDescriptor
from cupy_backends.cuda.libs.cutensor cimport ContractionDescriptor
from cupy_backends.cuda.libs.cutensor cimport ContractionFind
from cupy_backends.cuda.libs.cutensor cimport ContractionPlan

from cupy.core cimport core
from cupy.core cimport _reduction
from cupy.cuda cimport device
from cupy_backends.cuda.api cimport runtime
from cupy_backends.cuda.libs cimport cutensor


cdef dict _handles = {}
cdef dict _tensor_descriptors = {}
cdef dict _contraction_descriptors = {}
cdef dict _contraction_finds = {}
cdef dict _contraction_plans = {}
cdef dict _modes = {}
cdef dict _scalars = {}


cdef class Mode(object):

    cdef:
        object _array
        readonly int ndim
        readonly intptr_t data

    def __init__(self, mode):
        self._array = numpy.array(mode, dtype=numpy.int32)
        assert self._array.ndim == 1
        self.ndim = self._array.size
        self.data = self._array.ctypes.data

    def __repr__(self):
        return 'mode([' + ', '.join(self._array) + '])'


cdef class _Scalar(object):

    cdef:
        object _array
        readonly intptr_t ptr

    def __init__(self, value, dtype):
        self._array = numpy.asarray(value, dtype=dtype)
        self.ptr = self._array.ctypes.data

    def __repr__(self):
        return self._array.item()


cdef Handle _get_handle():
    cdef Handle handle
    cdef int dev = device.get_device_id()
    if dev not in _handles:
        handle = Handle()
        cutensor.init(handle)
        _handles[dev] = handle
        return handle
    return _handles[dev]


cdef int _get_cuda_dtype(numpy_dtype) except -1:
    if numpy_dtype == numpy.float16:
        return runtime.CUDA_R_16F
    elif numpy_dtype == numpy.float32:
        return runtime.CUDA_R_32F
    elif numpy_dtype == numpy.float64:
        return runtime.CUDA_R_64F
    elif numpy_dtype == numpy.complex64:
        return runtime.CUDA_C_32F
    elif numpy_dtype == numpy.complex128:
        return runtime.CUDA_C_64F
    else:
        raise TypeError('Dtype {} is not supported'.format(numpy_dtype))


cdef int _get_cutensor_dtype(numpy_dtype) except -1:
    if numpy_dtype == numpy.float16:
        return cutensor.R_MIN_16F
    elif numpy_dtype == numpy.float32:
        return cutensor.R_MIN_32F
    elif numpy_dtype == numpy.float64:
        return cutensor.R_MIN_64F
    elif numpy_dtype == numpy.complex64:
        return cutensor.C_MIN_32F
    elif numpy_dtype == numpy.complex128:
        return cutensor.C_MIN_64F
    else:
        raise TypeError('Dtype {} is not supported'.format(numpy_dtype))


def create_mode(*mode):
    """Create the tensor mode from the given integers or characters.

    Args:
        mode (tuple of int/str): A tuple that holds the labels of the modes
            of tensor A (e.g., if A_{x,y,z}, mode_A = {'x','y','z'})
    """
    integer_mode = []
    for x in mode:
        if isinstance(x, int):
            integer_mode.append(x)
        elif isinstance(x, str):
            integer_mode.append(ord(x))
        else:
            raise TypeError('Cannot create tensor mode: {}'.format(type(x)))
    return Mode(integer_mode)


cdef inline Mode _auto_create_mode(ndarray array, mode):
    if not isinstance(mode, Mode):
        mode = create_mode(*mode)
    if array.ndim != mode.ndim:
        raise ValueError(
            'ndim mismatch: {} != {}'.format(array.ndim, mode.ndim))
    return mode


cdef inline _set_compute_dtype(array_dtype, compute_dtype=None):
    if compute_dtype is None:
        if array_dtype == numpy.float16:
            compute_dtype = numpy.float32
        else:
            compute_dtype = array_dtype
    return compute_dtype


cdef inline _Scalar _create_scalar(scale, dtype):
    cdef _Scalar scalar
    key = (scale, dtype)
    if key in _scalars:
        scalar = _scalars[key]
    else:
        scalar = _Scalar(scale, dtype)
        _scalars[key] = scalar
    return scalar


cdef inline Mode _create_mode_with_cache(axis_or_ndim):
    cdef Mode mode
    if axis_or_ndim in _modes:
        mode = _modes[axis_or_ndim]
    else:
        if type(axis_or_ndim) is int:
            mode = Mode(tuple(range(axis_or_ndim)))
        else:
            mode = Mode(axis_or_ndim)
        _modes[axis_or_ndim] = mode
    return mode


cpdef TensorDescriptor create_tensor_descriptor(
        ndarray a, int uop=cutensor.OP_IDENTITY, Handle handle=None):
    """Create a tensor descriptor

    Args:
        a (cupy.ndarray): tensor for which a descritpor are created.
        uop (cutensorOperator_t): unary operator that will be applied to each
            element of the corresponding tensor in a lazy fashion (i.e., the
            algorithm uses this tensor as its operand only once). The
            original data of this tensor remains unchanged.

    Returns:
        (Descriptor): A instance of class Descriptor which holds a pointer to
            tensor descriptor and its destructor.
    """
    if handle is None:
        handle = _get_handle()
    key = (handle.ptr, a.dtype, tuple(a.shape), tuple(a.strides), uop)
    if key in _tensor_descriptors:
        desc = _tensor_descriptors[key]
        return desc
    num_modes = a.ndim
    extent = numpy.array(a.shape, dtype=numpy.int64)
    stride = numpy.array(a.strides, dtype=numpy.int64) // a.itemsize
    cuda_dtype = _get_cuda_dtype(a.dtype)
    desc = TensorDescriptor()
    cutensor.initTensorDescriptor(
        handle, desc, num_modes, extent.ctypes.data, stride.ctypes.data,
        cuda_dtype, uop)
    _tensor_descriptors[key] = desc
    return desc


def elementwise_trinary(
        alpha, ndarray A, TensorDescriptor desc_A, mode_A,
        beta, ndarray B, TensorDescriptor desc_B, mode_B,
        gamma, ndarray C, TensorDescriptor desc_C, mode_C,
        ndarray out=None,
        op_AB=cutensor.OP_ADD, op_ABC=cutensor.OP_ADD, compute_dtype=None):
    """Element-wise tensor operation for three input tensors

    This function performs a element-wise tensor operation of the form:

        D_{Pi^C(i_0,i_1,...,i_nc)} =
            op_ABC(op_AB(alpha * uop_A(A_{Pi^A(i_0,i_1,...,i_na)}),
                         beta  * uop_B(B_{Pi^B(i_0,i_1,...,i_nb)})),
                         gamma * uop_C(C_{Pi^C(i_0,i_1,...,i_nc)}))

    See cupy/cuda/cutensor.elementwiseTrinary() for details.

    Args:
        alpha (scalar or 0-dim numpy.ndarray): Scaling factor for tensor A.
        A (cupy.ndarray): Input tensor.
        desc_A (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor A.
        mode_A (cutensor.Mode): A mode object created by `create_mode`.
        beta (scalar or 0-dim numpy.ndarray): Scaling factor for tensor B.
        B (cupy.ndarray): Input tensor.
        desc_B (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor B.
        mode_B (cutensor.Mode): A mode object created by `create_mode`.
        gamma (scalar or 0-dim numpy.ndarray): Scaling factor for tensor C.
        C (cupy.ndarray): Input tensor.
        desc_C (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor C.
        mode_C (cutensor.Mode): A mode object created by `create_mode`.
        out (cupy.ndarray): Output tensor.
        op_AB (cutensorOperator_t): Element-wise binary operator.
        op_ABC (cutensorOperator_t): Element-wise binary operator.
        compute_dtype (numpy.dtype): Compute type for the intermediate
            computation.

    Returns:
        out (cupy.ndarray): Output tensor.

    Examples:
        See examples/cutensor/elementwise_trinary.py
    """
    if not (A.dtype == B.dtype == C.dtype):
        raise ValueError(
            'dtype mismatch: ({}, {}, {})'.format(A.dtype, B.dtype, C.dtype))
    if not (A._c_contiguous and B._c_contiguous and C._c_contiguous):
        raise ValueError('The inputs should be contiguous arrays.')

    if out is None:
        out = core._ndarray_init(C._shape, dtype=C.dtype)
    elif C.dtype != out.dtype:
        raise ValueError('dtype mismatch: {} != {}'.format(C.dtype, out.dtype))
    elif not internal.vector_equal(C._shape, out._shape):
        raise ValueError('shape mismatch: {} != {}'.format(C.shape, out.shape))
    elif not out._c_contiguous:
        raise ValueError('`out` should be a contiguous array.')

    if compute_dtype is None:
        compute_dtype = A.dtype

    return _elementwise_trinary_impl(
        _get_handle(),
        _create_scalar(alpha, compute_dtype),
        A, desc_A, _auto_create_mode(A, mode_A),
        _create_scalar(beta, compute_dtype),
        B, desc_B, _auto_create_mode(B, mode_B),
        _create_scalar(gamma, compute_dtype),
        C, desc_C, _auto_create_mode(C, mode_C),
        out, op_AB, op_ABC, _get_cuda_dtype(compute_dtype))


cdef inline ndarray _elementwise_trinary_impl(
        Handle handle,
        _Scalar alpha, ndarray A, TensorDescriptor desc_A, Mode mode_A,
        _Scalar beta, ndarray B, TensorDescriptor desc_B, Mode mode_B,
        _Scalar gamma, ndarray C, TensorDescriptor desc_C, Mode mode_C,
        ndarray out, int op_AB, int op_ABC, int compute_type):
    cutensor.elementwiseTrinary(
        handle,
        alpha.ptr, A.data.ptr, desc_A, mode_A.data,
        beta.ptr, B.data.ptr, desc_B, mode_B.data,
        gamma.ptr, C.data.ptr, desc_C, mode_C.data,
        out.data.ptr, desc_C, mode_C.data,
        op_AB, op_ABC, compute_type)
    return out


def elementwise_binary(
        alpha, ndarray A, TensorDescriptor desc_A, mode_A,
        gamma, ndarray C, TensorDescriptor desc_C, mode_C,
        ndarray out=None,
        op_AC=cutensor.OP_ADD, compute_dtype=None):
    """Element-wise tensor operation for two input tensors

    This function performs a element-wise tensor operation of the form:

        D_{Pi^C(i_0,i_1,...,i_n)} =
            op_AC(alpha * uop_A(A_{Pi^A(i_0,i_1,...,i_n)}),
                  gamma * uop_C(C_{Pi^C(i_0,i_1,...,i_n)}))

    See elementwise_trinary() for details.

    Examples:
        See examples/cutensor/elementwise_binary.py
    """
    if A.dtype != C.dtype:
        raise ValueError('dtype mismatch: {} != {}'.format(A.dtype, C.dtype))
    if not (A._c_contiguous and C._c_contiguous):
        raise ValueError('The inputs should be contiguous arrays.')

    if out is None:
        out = core._ndarray_init(C._shape, dtype=C.dtype)
    elif C.dtype != out.dtype:
        raise ValueError('dtype mismatch: {} != {}'.format(C.dtype, out.dtype))
    elif not internal.vector_equal(C._shape, out._shape):
        raise ValueError('shape mismatch: {} != {}'.format(C.shape, out.shape))
    elif not out._c_contiguous:
        raise ValueError('`out` should be a contiguous array.')

    if compute_dtype is None:
        compute_dtype = A.dtype

    return _elementwise_binary_impl(
        _get_handle(),
        _create_scalar(alpha, compute_dtype),
        A, desc_A, _auto_create_mode(A, mode_A),
        _create_scalar(gamma, compute_dtype),
        C, desc_C, _auto_create_mode(A, mode_C),
        out, op_AC, _get_cuda_dtype(compute_dtype)
    )


cdef inline ndarray _elementwise_binary_impl(
        Handle handle,
        _Scalar alpha, ndarray A, TensorDescriptor desc_A, Mode mode_A,
        _Scalar gamma, ndarray C, TensorDescriptor desc_C, Mode mode_C,
        ndarray out, int op_AC, int compute_type):
    cutensor.elementwiseBinary(
        handle,
        alpha.ptr, A.data.ptr, desc_A, mode_A.data,
        gamma.ptr, C.data.ptr, desc_C, mode_C.data,
        out.data.ptr, desc_C, mode_C.data,
        op_AC, compute_type)
    return out


cdef inline ContractionDescriptor _create_contraction_descriptor(
        Handle handle,
        ndarray A, TensorDescriptor desc_A, Mode mode_A,
        ndarray B, TensorDescriptor desc_B, Mode mode_B,
        ndarray C, TensorDescriptor desc_C, Mode mode_C,
        int cutensor_compute_type):
    """Create a contraction descriptor"""
    cdef uint32_t alignment_req_A = cutensor.getAlignmentRequirement(
        handle, A.data.ptr, desc_A)
    cdef uint32_t alignment_req_B = cutensor.getAlignmentRequirement(
        handle, B.data.ptr, desc_B)
    cdef uint32_t alignment_req_C = cutensor.getAlignmentRequirement(
        handle, C.data.ptr, desc_C)
    cdef ContractionDescriptor desc

    key = (handle.ptr, cutensor_compute_type,
           desc_A.ptr, mode_A.data, alignment_req_A,
           desc_B.ptr, mode_B.data, alignment_req_B,
           desc_C.ptr, mode_C.data, alignment_req_C)
    if key in _contraction_descriptors:
        desc = _contraction_descriptors[key]
        return desc

    desc = ContractionDescriptor()
    cutensor.initContractionDescriptor(
        handle,
        desc,
        desc_A, mode_A.data, alignment_req_A,
        desc_B, mode_B.data, alignment_req_B,
        desc_C, mode_C.data, alignment_req_C,
        desc_C, mode_C.data, alignment_req_C,
        cutensor_compute_type)
    _contraction_descriptors[key] = desc
    return desc


cdef inline ContractionFind _create_contraction_find(Handle handle, int algo):
    """Create a contraction find"""
    cdef ContractionFind find

    key = (handle.ptr, algo)
    if key in _contraction_finds:
        find = _contraction_finds[key]
    else:
        find = ContractionFind()
        cutensor.initContractionFind(handle, find, algo)
        _contraction_finds[key] = find
    return find


cdef inline ContractionPlan _create_contraction_plan(
        Handle handle,
        ContractionDescriptor desc, ContractionFind find, uint64_t ws_size):
    """Create a contraction plan"""
    cdef ContractionPlan plan

    key = (handle.ptr, desc.ptr, find.ptr, ws_size)
    if key in _contraction_plans:
        plan = _contraction_plans[key]
    else:
        plan = ContractionPlan()
        cutensor.initContractionPlan(handle, plan, desc, find, ws_size)
        _contraction_plans[key] = plan
    return plan


def contraction(
        alpha, ndarray A, TensorDescriptor desc_A, mode_A,
        ndarray B, TensorDescriptor desc_B, mode_B,
        beta, ndarray C, TensorDescriptor desc_C, mode_C,
        compute_dtype=None,
        int algo=cutensor.ALGO_DEFAULT,
        int ws_pref=cutensor.WORKSPACE_RECOMMENDED):
    """General tensor contraction

    This routine computes the tensor contraction:

        C = alpha * uop_A(A) * uop_B(B) + beta * uop_C(C)

    See cupy/cuda/cutensor.contraction for details.

    Args:
        alpha (scalar or 0-dim numpy.ndarray): Scaling factor for A * B.
        A (cupy.ndarray): Input tensor.
        desc_A (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor A.
        mode_A (cutensor.Mode): A mode object created by `create_mode`.
        B (cupy.ndarray): Input tensor.
        desc_B (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor B.
        mode_B (cutensor.Mode): A mode object created by `create_mode`.
        beta (scalar or 0-dim numpy.ndarray): Scaling factor for C.
        C (cupy.ndarray): Input/output tensor.
        desc_C (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor C.
        mode_C (cutensor.Mode): A mode object created by `create_mode`.
        compute_dtype (numpy.dtype): Compute type for the intermediate
            computation.
        algo (cutensorAlgo_t): Allows users to select a specific algorithm.
            ALGO_DEFAULT lets the heuristic choose the algorithm.
            Any value >= 0 selects a specific GEMM-like algorithm and
            deactivates the heuristic. If a specified algorithm is not
            supported, STATUS_NOT_SUPPORTED is returned.
        ws_pref (cutensorWorksizePreference_t): User preference for the
            workspace of cuTensor.

    Returns:
        out (cupy.ndarray): Output tensor.

    Examples:
        See examples/cutensor/contraction.py
    """
    if not (A.dtype == B.dtype == C.dtype):
        raise ValueError(
            'dtype mismatch: ({}, {}, {})'.format(A.dtype, B.dtype, C.dtype))
    if not (A._c_contiguous and B._c_contiguous and C._c_contiguous):
        raise ValueError('The inputs should be contiguous arrays.')

    compute_dtype = _set_compute_dtype(A.dtype, compute_dtype)

    return _contraction_impl(
        _get_handle(),
        _create_scalar(alpha, compute_dtype),
        A, desc_A, _auto_create_mode(A, mode_A),
        B, desc_B, _auto_create_mode(B, mode_B),
        _create_scalar(beta, compute_dtype),
        C, desc_C, _auto_create_mode(C, mode_C),
        _get_cutensor_dtype(compute_dtype), algo, ws_pref)


cdef inline ndarray _contraction_impl(
        Handle handle,
        _Scalar alpha, ndarray A, TensorDescriptor desc_A, Mode mode_A,
        ndarray B, TensorDescriptor desc_B, Mode mode_B,
        _Scalar beta, ndarray C, TensorDescriptor desc_C, Mode mode_C,
        int cutensor_compute_type, int algo, int ws_pref):
    cdef ContractionDescriptor desc
    cdef ContractionFind find
    cdef ContractionPlan plan
    cdef uint64_t ws_size
    cdef ndarray out, ws

    out = C

    desc = _create_contraction_descriptor(
        handle,
        A, desc_A, mode_A,
        B, desc_B, mode_B,
        C, desc_C, mode_C,
        cutensor_compute_type)

    find = _create_contraction_find(handle, algo)

    # Allocate workspace
    ws_size = cutensor.contractionGetWorkspace(handle, desc, find, ws_pref)
    try:
        ws = core._ndarray_init(shape_t(1, ws_size), dtype=numpy.int8)
    except Exception:
        warnings.warn('cuTENSOR: failed to allocate memory of workspace '
                      'with preference ({}) and size ({}).'
                      ''.format(ws_pref, ws_size))
        ws_size = cutensor.contractionGetWorkspace(
            handle, desc, find, cutensor.WORKSPACE_MIN)
        ws = core._ndarray_init(shape_t(1, ws_size), dtype=numpy.int8)

    plan = _create_contraction_plan(handle, desc, find, ws_size)

    cutensor.contraction(
        handle, plan,
        alpha.ptr, A.data.ptr, B.data.ptr,
        beta.ptr, C.data.ptr, out.data.ptr,
        ws.data.ptr, ws_size)
    return out


def contraction_max_algos():
    """Returns the maximum number of algorithms for cutensor()

    See cupy/cuda/cutensor.contractionMaxAlgos() for details.
    """
    return cutensor.contractionMaxAlgos()


def reduction(
        alpha, ndarray A, TensorDescriptor desc_A, mode_A,
        beta, ndarray C, TensorDescriptor desc_C, mode_C,
        int reduce_op=cutensor.OP_ADD, compute_dtype=None):
    """Tensor reduction

    This routine computes the tensor reduction:

        C = alpha * reduce_op(uop_A(A)) + beta * uop_C(C))

    See :func:`cupy.cuda.cutensor.reduction` for details.

    Args:
        alpha (scalar or 0-dim numpy.ndarray): Scaling factor for A.
        A (cupy.ndarray): Input tensor.
        desc_A (class Descriptor): A descriptor that holds the information
            about the data type, modes, strides and unary operator (uop_A) of
            tensor A.
        mode_A (cutensor.Mode): A mode object created by `create_mode`.
        beta (scalar or 0-dim numpy.ndarray): Scaling factor for C.
        C (cupy.ndarray): Input/output tensor.
        desc_C (class Descriptor): A descriptor that holds the information
            about the data type, modes, strides and unary operator (uop_C) of
            tensor C.
        mode_C (cutensor.Mode): A mode object created by `create_mode`.
        reduce_op (cutensorOperator_t): Binary operator used to reduce A.
        compute_dtype (numpy.dtype): Compute type for the intermediate
            computation.

    Returns:
        out (cupy.ndarray): Output tensor.

    Examples:
        See examples/cutensor/reduction.py
    """
    cdef Handle handle

    if A.dtype != C.dtype:
        raise ValueError('dtype mismatch: {} != {}'.format(A.dtype, C.dtype))<commit_msg>Fix documentation<commit_after>import numpy
import warnings

import cupy

from libc.stdint cimport intptr_t, uint32_t, uint64_t
from cupy.core._carray cimport shape_t
from cupy.core.core cimport ndarray
from cupy.core cimport internal
from cupy_backends.cuda.libs.cutensor cimport Handle
from cupy_backends.cuda.libs.cutensor cimport TensorDescriptor
from cupy_backends.cuda.libs.cutensor cimport ContractionDescriptor
from cupy_backends.cuda.libs.cutensor cimport ContractionFind
from cupy_backends.cuda.libs.cutensor cimport ContractionPlan

from cupy.core cimport core
from cupy.core cimport _reduction
from cupy.cuda cimport device
from cupy_backends.cuda.api cimport runtime
from cupy_backends.cuda.libs cimport cutensor


cdef dict _handles = {}
cdef dict _tensor_descriptors = {}
cdef dict _contraction_descriptors = {}
cdef dict _contraction_finds = {}
cdef dict _contraction_plans = {}
cdef dict _modes = {}
cdef dict _scalars = {}


cdef class Mode(object):

    cdef:
        object _array
        readonly int ndim
        readonly intptr_t data

    def __init__(self, mode):
        self._array = numpy.array(mode, dtype=numpy.int32)
        assert self._array.ndim == 1
        self.ndim = self._array.size
        self.data = self._array.ctypes.data

    def __repr__(self):
        return 'mode([' + ', '.join(self._array) + '])'


cdef class _Scalar(object):

    cdef:
        object _array
        readonly intptr_t ptr

    def __init__(self, value, dtype):
        self._array = numpy.asarray(value, dtype=dtype)
        self.ptr = self._array.ctypes.data

    def __repr__(self):
        return self._array.item()


cdef Handle _get_handle():
    cdef Handle handle
    cdef int dev = device.get_device_id()
    if dev not in _handles:
        handle = Handle()
        cutensor.init(handle)
        _handles[dev] = handle
        return handle
    return _handles[dev]


cdef int _get_cuda_dtype(numpy_dtype) except -1:
    if numpy_dtype == numpy.float16:
        return runtime.CUDA_R_16F
    elif numpy_dtype == numpy.float32:
        return runtime.CUDA_R_32F
    elif numpy_dtype == numpy.float64:
        return runtime.CUDA_R_64F
    elif numpy_dtype == numpy.complex64:
        return runtime.CUDA_C_32F
    elif numpy_dtype == numpy.complex128:
        return runtime.CUDA_C_64F
    else:
        raise TypeError('Dtype {} is not supported'.format(numpy_dtype))


cdef int _get_cutensor_dtype(numpy_dtype) except -1:
    if numpy_dtype == numpy.float16:
        return cutensor.R_MIN_16F
    elif numpy_dtype == numpy.float32:
        return cutensor.R_MIN_32F
    elif numpy_dtype == numpy.float64:
        return cutensor.R_MIN_64F
    elif numpy_dtype == numpy.complex64:
        return cutensor.C_MIN_32F
    elif numpy_dtype == numpy.complex128:
        return cutensor.C_MIN_64F
    else:
        raise TypeError('Dtype {} is not supported'.format(numpy_dtype))


def create_mode(*mode):
    """Create the tensor mode from the given integers or characters.

    Args:
        mode (tuple of int/str): A tuple that holds the labels of the modes
            of tensor A (e.g., if A_{x,y,z}, mode_A = {'x','y','z'})
    """
    integer_mode = []
    for x in mode:
        if isinstance(x, int):
            integer_mode.append(x)
        elif isinstance(x, str):
            integer_mode.append(ord(x))
        else:
            raise TypeError('Cannot create tensor mode: {}'.format(type(x)))
    return Mode(integer_mode)


cdef inline Mode _auto_create_mode(ndarray array, mode):
    if not isinstance(mode, Mode):
        mode = create_mode(*mode)
    if array.ndim != mode.ndim:
        raise ValueError(
            'ndim mismatch: {} != {}'.format(array.ndim, mode.ndim))
    return mode


cdef inline _set_compute_dtype(array_dtype, compute_dtype=None):
    if compute_dtype is None:
        if array_dtype == numpy.float16:
            compute_dtype = numpy.float32
        else:
            compute_dtype = array_dtype
    return compute_dtype


cdef inline _Scalar _create_scalar(scale, dtype):
    cdef _Scalar scalar
    key = (scale, dtype)
    if key in _scalars:
        scalar = _scalars[key]
    else:
        scalar = _Scalar(scale, dtype)
        _scalars[key] = scalar
    return scalar


cdef inline Mode _create_mode_with_cache(axis_or_ndim):
    cdef Mode mode
    if axis_or_ndim in _modes:
        mode = _modes[axis_or_ndim]
    else:
        if type(axis_or_ndim) is int:
            mode = Mode(tuple(range(axis_or_ndim)))
        else:
            mode = Mode(axis_or_ndim)
        _modes[axis_or_ndim] = mode
    return mode


cpdef TensorDescriptor create_tensor_descriptor(
        ndarray a, int uop=cutensor.OP_IDENTITY, Handle handle=None):
    """Create a tensor descriptor

    Args:
        a (cupy.ndarray): tensor for which a descritpor are created.
        uop (cutensorOperator_t): unary operator that will be applied to each
            element of the corresponding tensor in a lazy fashion (i.e., the
            algorithm uses this tensor as its operand only once). The
            original data of this tensor remains unchanged.

    Returns:
        (Descriptor): A instance of class Descriptor which holds a pointer to
            tensor descriptor and its destructor.
    """
    if handle is None:
        handle = _get_handle()
    key = (handle.ptr, a.dtype, tuple(a.shape), tuple(a.strides), uop)
    if key in _tensor_descriptors:
        desc = _tensor_descriptors[key]
        return desc
    num_modes = a.ndim
    extent = numpy.array(a.shape, dtype=numpy.int64)
    stride = numpy.array(a.strides, dtype=numpy.int64) // a.itemsize
    cuda_dtype = _get_cuda_dtype(a.dtype)
    desc = TensorDescriptor()
    cutensor.initTensorDescriptor(
        handle, desc, num_modes, extent.ctypes.data, stride.ctypes.data,
        cuda_dtype, uop)
    _tensor_descriptors[key] = desc
    return desc


def elementwise_trinary(
        alpha, ndarray A, TensorDescriptor desc_A, mode_A,
        beta, ndarray B, TensorDescriptor desc_B, mode_B,
        gamma, ndarray C, TensorDescriptor desc_C, mode_C,
        ndarray out=None,
        op_AB=cutensor.OP_ADD, op_ABC=cutensor.OP_ADD, compute_dtype=None):
    """Element-wise tensor operation for three input tensors

    This function performs a element-wise tensor operation of the form:

        D_{Pi^C(i_0,i_1,...,i_nc)} =
            op_ABC(op_AB(alpha * uop_A(A_{Pi^A(i_0,i_1,...,i_na)}),
                         beta  * uop_B(B_{Pi^B(i_0,i_1,...,i_nb)})),
                         gamma * uop_C(C_{Pi^C(i_0,i_1,...,i_nc)}))

    See cupy/cuda/cutensor.elementwiseTrinary() for details.

    Args:
        alpha (scalar): Scaling factor for tensor A.
        A (cupy.ndarray): Input tensor.
        desc_A (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor A.
        mode_A (cutensor.Mode): A mode object created by `create_mode`.
        beta (scalar): Scaling factor for tensor B.
        B (cupy.ndarray): Input tensor.
        desc_B (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor B.
        mode_B (cutensor.Mode): A mode object created by `create_mode`.
        gamma (scalar): Scaling factor for tensor C.
        C (cupy.ndarray): Input tensor.
        desc_C (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor C.
        mode_C (cutensor.Mode): A mode object created by `create_mode`.
        out (cupy.ndarray): Output tensor.
        op_AB (cutensorOperator_t): Element-wise binary operator.
        op_ABC (cutensorOperator_t): Element-wise binary operator.
        compute_dtype (numpy.dtype): Compute type for the intermediate
            computation.

    Returns:
        out (cupy.ndarray): Output tensor.

    Examples:
        See examples/cutensor/elementwise_trinary.py
    """
    if not (A.dtype == B.dtype == C.dtype):
        raise ValueError(
            'dtype mismatch: ({}, {}, {})'.format(A.dtype, B.dtype, C.dtype))
    if not (A._c_contiguous and B._c_contiguous and C._c_contiguous):
        raise ValueError('The inputs should be contiguous arrays.')

    if out is None:
        out = core._ndarray_init(C._shape, dtype=C.dtype)
    elif C.dtype != out.dtype:
        raise ValueError('dtype mismatch: {} != {}'.format(C.dtype, out.dtype))
    elif not internal.vector_equal(C._shape, out._shape):
        raise ValueError('shape mismatch: {} != {}'.format(C.shape, out.shape))
    elif not out._c_contiguous:
        raise ValueError('`out` should be a contiguous array.')

    if compute_dtype is None:
        compute_dtype = A.dtype

    return _elementwise_trinary_impl(
        _get_handle(),
        _create_scalar(alpha, compute_dtype),
        A, desc_A, _auto_create_mode(A, mode_A),
        _create_scalar(beta, compute_dtype),
        B, desc_B, _auto_create_mode(B, mode_B),
        _create_scalar(gamma, compute_dtype),
        C, desc_C, _auto_create_mode(C, mode_C),
        out, op_AB, op_ABC, _get_cuda_dtype(compute_dtype))


cdef inline ndarray _elementwise_trinary_impl(
        Handle handle,
        _Scalar alpha, ndarray A, TensorDescriptor desc_A, Mode mode_A,
        _Scalar beta, ndarray B, TensorDescriptor desc_B, Mode mode_B,
        _Scalar gamma, ndarray C, TensorDescriptor desc_C, Mode mode_C,
        ndarray out, int op_AB, int op_ABC, int compute_type):
    cutensor.elementwiseTrinary(
        handle,
        alpha.ptr, A.data.ptr, desc_A, mode_A.data,
        beta.ptr, B.data.ptr, desc_B, mode_B.data,
        gamma.ptr, C.data.ptr, desc_C, mode_C.data,
        out.data.ptr, desc_C, mode_C.data,
        op_AB, op_ABC, compute_type)
    return out


def elementwise_binary(
        alpha, ndarray A, TensorDescriptor desc_A, mode_A,
        gamma, ndarray C, TensorDescriptor desc_C, mode_C,
        ndarray out=None,
        op_AC=cutensor.OP_ADD, compute_dtype=None):
    """Element-wise tensor operation for two input tensors

    This function performs a element-wise tensor operation of the form:

        D_{Pi^C(i_0,i_1,...,i_n)} =
            op_AC(alpha * uop_A(A_{Pi^A(i_0,i_1,...,i_n)}),
                  gamma * uop_C(C_{Pi^C(i_0,i_1,...,i_n)}))

    See elementwise_trinary() for details.

    Examples:
        See examples/cutensor/elementwise_binary.py
    """
    if A.dtype != C.dtype:
        raise ValueError('dtype mismatch: {} != {}'.format(A.dtype, C.dtype))
    if not (A._c_contiguous and C._c_contiguous):
        raise ValueError('The inputs should be contiguous arrays.')

    if out is None:
        out = core._ndarray_init(C._shape, dtype=C.dtype)
    elif C.dtype != out.dtype:
        raise ValueError('dtype mismatch: {} != {}'.format(C.dtype, out.dtype))
    elif not internal.vector_equal(C._shape, out._shape):
        raise ValueError('shape mismatch: {} != {}'.format(C.shape, out.shape))
    elif not out._c_contiguous:
        raise ValueError('`out` should be a contiguous array.')

    if compute_dtype is None:
        compute_dtype = A.dtype

    return _elementwise_binary_impl(
        _get_handle(),
        _create_scalar(alpha, compute_dtype),
        A, desc_A, _auto_create_mode(A, mode_A),
        _create_scalar(gamma, compute_dtype),
        C, desc_C, _auto_create_mode(A, mode_C),
        out, op_AC, _get_cuda_dtype(compute_dtype)
    )


cdef inline ndarray _elementwise_binary_impl(
        Handle handle,
        _Scalar alpha, ndarray A, TensorDescriptor desc_A, Mode mode_A,
        _Scalar gamma, ndarray C, TensorDescriptor desc_C, Mode mode_C,
        ndarray out, int op_AC, int compute_type):
    cutensor.elementwiseBinary(
        handle,
        alpha.ptr, A.data.ptr, desc_A, mode_A.data,
        gamma.ptr, C.data.ptr, desc_C, mode_C.data,
        out.data.ptr, desc_C, mode_C.data,
        op_AC, compute_type)
    return out


cdef inline ContractionDescriptor _create_contraction_descriptor(
        Handle handle,
        ndarray A, TensorDescriptor desc_A, Mode mode_A,
        ndarray B, TensorDescriptor desc_B, Mode mode_B,
        ndarray C, TensorDescriptor desc_C, Mode mode_C,
        int cutensor_compute_type):
    """Create a contraction descriptor"""
    cdef uint32_t alignment_req_A = cutensor.getAlignmentRequirement(
        handle, A.data.ptr, desc_A)
    cdef uint32_t alignment_req_B = cutensor.getAlignmentRequirement(
        handle, B.data.ptr, desc_B)
    cdef uint32_t alignment_req_C = cutensor.getAlignmentRequirement(
        handle, C.data.ptr, desc_C)
    cdef ContractionDescriptor desc

    key = (handle.ptr, cutensor_compute_type,
           desc_A.ptr, mode_A.data, alignment_req_A,
           desc_B.ptr, mode_B.data, alignment_req_B,
           desc_C.ptr, mode_C.data, alignment_req_C)
    if key in _contraction_descriptors:
        desc = _contraction_descriptors[key]
        return desc

    desc = ContractionDescriptor()
    cutensor.initContractionDescriptor(
        handle,
        desc,
        desc_A, mode_A.data, alignment_req_A,
        desc_B, mode_B.data, alignment_req_B,
        desc_C, mode_C.data, alignment_req_C,
        desc_C, mode_C.data, alignment_req_C,
        cutensor_compute_type)
    _contraction_descriptors[key] = desc
    return desc


cdef inline ContractionFind _create_contraction_find(Handle handle, int algo):
    """Create a contraction find"""
    cdef ContractionFind find

    key = (handle.ptr, algo)
    if key in _contraction_finds:
        find = _contraction_finds[key]
    else:
        find = ContractionFind()
        cutensor.initContractionFind(handle, find, algo)
        _contraction_finds[key] = find
    return find


cdef inline ContractionPlan _create_contraction_plan(
        Handle handle,
        ContractionDescriptor desc, ContractionFind find, uint64_t ws_size):
    """Create a contraction plan"""
    cdef ContractionPlan plan

    key = (handle.ptr, desc.ptr, find.ptr, ws_size)
    if key in _contraction_plans:
        plan = _contraction_plans[key]
    else:
        plan = ContractionPlan()
        cutensor.initContractionPlan(handle, plan, desc, find, ws_size)
        _contraction_plans[key] = plan
    return plan


def contraction(
        alpha, ndarray A, TensorDescriptor desc_A, mode_A,
        ndarray B, TensorDescriptor desc_B, mode_B,
        beta, ndarray C, TensorDescriptor desc_C, mode_C,
        compute_dtype=None,
        int algo=cutensor.ALGO_DEFAULT,
        int ws_pref=cutensor.WORKSPACE_RECOMMENDED):
    """General tensor contraction

    This routine computes the tensor contraction:

        C = alpha * uop_A(A) * uop_B(B) + beta * uop_C(C)

    See cupy/cuda/cutensor.contraction for details.

    Args:
        alpha (scalar): Scaling factor for A * B.
        A (cupy.ndarray): Input tensor.
        desc_A (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor A.
        mode_A (cutensor.Mode): A mode object created by `create_mode`.
        B (cupy.ndarray): Input tensor.
        desc_B (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor B.
        mode_B (cutensor.Mode): A mode object created by `create_mode`.
        beta (scalar): Scaling factor for C.
        C (cupy.ndarray): Input/output tensor.
        desc_C (class Descriptor): A descriptor that holds the information
            about the data type, modes, and strides of tensor C.
        mode_C (cutensor.Mode): A mode object created by `create_mode`.
        compute_dtype (numpy.dtype): Compute type for the intermediate
            computation.
        algo (cutensorAlgo_t): Allows users to select a specific algorithm.
            ALGO_DEFAULT lets the heuristic choose the algorithm.
            Any value >= 0 selects a specific GEMM-like algorithm and
            deactivates the heuristic. If a specified algorithm is not
            supported, STATUS_NOT_SUPPORTED is returned.
        ws_pref (cutensorWorksizePreference_t): User preference for the
            workspace of cuTensor.

    Returns:
        out (cupy.ndarray): Output tensor.

    Examples:
        See examples/cutensor/contraction.py
    """
    if not (A.dtype == B.dtype == C.dtype):
        raise ValueError(
            'dtype mismatch: ({}, {}, {})'.format(A.dtype, B.dtype, C.dtype))
    if not (A._c_contiguous and B._c_contiguous and C._c_contiguous):
        raise ValueError('The inputs should be contiguous arrays.')

    compute_dtype = _set_compute_dtype(A.dtype, compute_dtype)

    return _contraction_impl(
        _get_handle(),
        _create_scalar(alpha, compute_dtype),
        A, desc_A, _auto_create_mode(A, mode_A),
        B, desc_B, _auto_create_mode(B, mode_B),
        _create_scalar(beta, compute_dtype),
        C, desc_C, _auto_create_mode(C, mode_C),
        _get_cutensor_dtype(compute_dtype), algo, ws_pref)


cdef inline ndarray _contraction_impl(
        Handle handle,
        _Scalar alpha, ndarray A, TensorDescriptor desc_A, Mode mode_A,
        ndarray B, TensorDescriptor desc_B, Mode mode_B,
        _Scalar beta, ndarray C, TensorDescriptor desc_C, Mode mode_C,
        int cutensor_compute_type, int algo, int ws_pref):
    cdef ContractionDescriptor desc
    cdef ContractionFind find
    cdef ContractionPlan plan
    cdef uint64_t ws_size
    cdef ndarray out, ws

    out = C

    desc = _create_contraction_descriptor(
        handle,
        A, desc_A, mode_A,
        B, desc_B, mode_B,
        C, desc_C, mode_C,
        cutensor_compute_type)

    find = _create_contraction_find(handle, algo)

    # Allocate workspace
    ws_size = cutensor.contractionGetWorkspace(handle, desc, find, ws_pref)
    try:
        ws = core._ndarray_init(shape_t(1, ws_size), dtype=numpy.int8)
    except Exception:
        warnings.warn('cuTENSOR: failed to allocate memory of workspace '
                      'with preference ({}) and size ({}).'
                      ''.format(ws_pref, ws_size))
        ws_size = cutensor.contractionGetWorkspace(
            handle, desc, find, cutensor.WORKSPACE_MIN)
        ws = core._ndarray_init(shape_t(1, ws_size), dtype=numpy.int8)

    plan = _create_contraction_plan(handle, desc, find, ws_size)

    cutensor.contraction(
        handle, plan,
        alpha.ptr, A.data.ptr, B.data.ptr,
        beta.ptr, C.data.ptr, out.data.ptr,
        ws.data.ptr, ws_size)
    return out


def contraction_max_algos():
    """Returns the maximum number of algorithms for cutensor()

    See cupy/cuda/cutensor.contractionMaxAlgos() for details.
    """
    return cutensor.contractionMaxAlgos()


def reduction(
        alpha, ndarray A, TensorDescriptor desc_A, mode_A,
        beta, ndarray C, TensorDescriptor desc_C, mode_C,
        int reduce_op=cutensor.OP_ADD, compute_dtype=None):
    """Tensor reduction

    This routine computes the tensor reduction:

        C = alpha * reduce_op(uop_A(A)) + beta * uop_C(C))

    See :func:`cupy.cuda.cutensor.reduction` for details.

    Args:
        alpha (scalar): Scaling factor for A.
        A (cupy.ndarray): Input tensor.
        desc_A (class Descriptor): A descriptor that holds the information
            about the data type, modes, strides and unary operator (uop_A) of
            tensor A.
        mode_A (cutensor.Mode): A mode object created by `create_mode`.
        beta (scalar): Scaling factor for C.
        C (cupy.ndarray): Input/output tensor.
        desc_C (class Descriptor): A descriptor that holds the information
            about the data type, modes, strides and unary operator (uop_C) of
            tensor C.
        mode_C (cutensor.Mode): A mode object created by `create_mode`.
        reduce_op (cutensorOperator_t): Binary operator used to reduce A.
        compute_dtype (numpy.dtype): Compute type for the intermediate
            computation.

    Returns:
        out (cupy.ndarray): Output tensor.

    Examples:
        See examples/cutensor/reduction.py
    """
    cdef Handle handle

    if A.dtype != C.dtype:
        raise ValueError('dtype mismatch: {} != {}'.format(A.dtype, C.dtype))
{'fraction_non_alphanumeric': 0.08941907179651493, 'fraction_numerical': 0.004510388682128509, 'mean_word_length': 2.755808571693048, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3990011', 'n_tokens_mistral': 12931, 'n_tokens_neox': 12713, 'n_words': 3517}