#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name(s) of the copyright holder(s) nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

def normalize_subscript(subscript):
    if '->' in subscript:
        subscript = subscript.split('->')
        lhs = subscript[0]
        rhs = subscript[1]
    else:
        lhs = subscript
        rhs = ''.join(sorted([s for s in set(subscript) if s != ',' and subscript.count(s) == 1]))
    if '...' in lhs:
        raise RuntimeError('Elipsis is currently unsupported')
    return lhs + '->' + rhs, ',' in lhs


# ! /usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name(s) of the copyright holder(s) nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from cutensor.c_extensions_utils import CustomExtension

einsum_torch = CustomExtension.Torch('cutensor.torch.binding',
                                     sources=['cutensor/torch/einsum.cc'])

einsum_tf = CustomExtension.Tensorflow(
    'cutensor.tensorflow.binding',
    sources=[
        'cutensor/tensorflow/einsum_kernel.cc',
        'cutensor/tensorflow/einsum_ops.cc',
        'cutensor/tensorflow/einsum_module.cc'
    ])


# ! /usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name(s) of the copyright holder(s) nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from setuptools import Extension
from distutils.spawn import find_executable
import os
import subprocess
import re

__all__ = ['CustomExtension']

include_dirs = []
library_dirs = []

cuda_nvcc = find_executable('nvcc')
cuda_root = os.path.join(os.path.dirname(cuda_nvcc), os.pardir)
cuda_version = re.search(
    r'release ([^,]*),',
    subprocess.check_output([cuda_nvcc, '--version']).decode('utf-8')).group(1)
include_dirs.append(os.path.join(cuda_root, 'include'))
library_dirs.append(os.path.join(cuda_root, 'lib64'))

if 'CUTENSOR_ROOT' in os.environ:
    root = os.environ['CUTENSOR_ROOT']
    include_dirs.append(os.path.join(root, 'include'))
    library_dirs.append(os.path.join(root, 'lib'))
    library_dirs.append(os.path.join(root, 'build/lib'))
    versioned_path = os.path.join(root, 'lib', cuda_version)
    if not os.path.exists(versioned_path):
        versioned_path = os.path.join(root, 'lib', cuda_version.split('.')[0])
    library_dirs.append(versioned_path)

class CustomExtension:
    modules = []

    @classmethod
    def Torch(cls, name, sources):
        try:
            import torch
            from torch.utils.cpp_extension import CUDAExtension
            ext = CUDAExtension(name,
                                sources=sources,
                                libraries=['cutensor'],
                                define_macros=[
                                    ('TORCH_API_INCLUDE_EXTENSION_H',),
                                    ('TORCH_EXTENSION_NAME',
                                     name.split('.')[-1]),
                                    ('_GLIBCXX_USE_CXX11_ABI',
                                     str(int(torch._C._GLIBCXX_USE_CXX11_ABI)))
                                ],
                                extra_compile_args=['-std=c++17', '-fopenmp'],
                                extra_link_args=['-std=c++17', '-fopenmp'],
                                include_dirs=include_dirs,
                                library_dirs=library_dirs,
                                runtime_library_dirs=library_dirs)
            cls.modules.append(ext)
            return ext
        except ImportError:
            return None

    @classmethod
    def Tensorflow(cls, name, sources):
        try:
            import tensorflow as tf
            ext = Extension(name,
                            sources=sources,
                            libraries=['cutensor', 'cudart'],
                            extra_compile_args=tf.sysconfig.get_compile_flags(),
                            extra_link_args=tf.sysconfig.get_link_flags() +
                            tf.sysconfig.get_compile_flags(),
                            define_macros=[('GOOGLE_CUDA', '1')],
                            include_dirs=include_dirs,
                            library_dirs=library_dirs,
                            runtime_library_dirs=library_dirs)
            cls.modules.append(ext)
        except ImportError:
            return None


CustomExtension

cuda_nvcc = find_executable('nvcc')
cuda_root = os.path.join(os.path.dirname(cuda_nvcc), os.pardir)
cuda_version = re.search(
    r'release ([^,]*),',
    subprocess.check_output([cuda_nvcc, '--version']).decode('utf-8')).group(1)
include_dirs.append(os.path.join(cuda_root, 'include'))
library_dirs.append(os.path.join(cuda_root, 'lib64'))

if 'CUTENSOR_ROOT' in os.environ:
    root = os.environ['CUTENSOR_ROOT']
    include_dirs.append(os.path.join(root, 'include'))
    library_dirs.append(os.path.join(root, 'lib'))
    library_dirs.append(os.path.join(root, 'build/lib'))
    versioned_path = os.path.join(root, 'lib', cuda_version)
    if not os.path.exists(versioned_path):
        versioned_path = os.path.join(root, 'lib', cuda_version.split('.')[0])
    library_dirs.append(versioned_path)

class CustomExtension:
    modules = []

    @classmethod
    def Torch(cls, name, sources):
        try:
            import torch
            from torch.utils.cpp_extension import CUDAExtension
            ext = CUDAExtension(name,
                                sources=sources,
                                libraries=['cutensor'],
                                define_macros=[
                                    ('TORCH_API_INCLUDE_EXTENSION_H',),
                                    ('TORCH_EXTENSION_NAME',
                                     name.split('.')[-1]),
                                    ('_GLIBCXX_USE_CXX11_ABI',
                                     str(int(torch._C._GLIBCXX_USE_CXX11_ABI)))
                                ],
                                extra_compile_args=['-std=c++17', '-fopenmp'],
                                extra_link_args=['-std=c++17', '-fopenmp'],
                                include_dirs=include_dirs,
                                library_dirs=library_dirs,
                                runtime_library_dirs=library_dirs)
            cls.modules.append(ext)
            return ext
        except ImportError:
            return None

    @classmethod
    def Tensorflow(cls, name, sources):
        try:
            import tensorflow as tf
            ext = Extension(name,
                            sources=sources,
                            libraries=['cutensor', 'cudart'],
                            extra_compile_args=tf.sysconfig.get_compile_flags(),
                            extra_link_args=tf.sysconfig.get_link_flags() +
                            tf.sysconfig.get_compile_flags(),
                            define_macros=[('GOOGLE_CUDA', '1')],
                            include_dirs=include_dirs,
                            library_dirs=library_dirs,
                            runtime_library_dirs=library_dirs)
            cls.modules.append(ext)
        except ImportError:
            return None


Torch

    r'release ([^,]*),',
    subprocess.check_output([cuda_nvcc, '--version']).decode('utf-8')).group(1)
include_dirs.append(os.path.join(cuda_root, 'include'))
library_dirs.append(os.path.join(cuda_root, 'lib64'))

if 'CUTENSOR_ROOT' in os.environ:
    root = os.environ['CUTENSOR_ROOT']
    include_dirs.append(os.path.join(root, 'include'))
    library_dirs.append(os.path.join(root, 'lib'))
    library_dirs.append(os.path.join(root, 'build/lib'))
    versioned_path = os.path.join(root, 'lib', cuda_version)
    if not os.path.exists(versioned_path):
        versioned_path = os.path.join(root, 'lib', cuda_version.split('.')[0])
    library_dirs.append(versioned_path)

class CustomExtension:
    modules = []

    @classmethod
    def Torch(cls, name, sources):
        try:
            import torch
            from torch.utils.cpp_extension import CUDAExtension
            ext = CUDAExtension(name,
                                sources=sources,
                                libraries=['cutensor'],
                                define_macros=[
                                    ('TORCH_API_INCLUDE_EXTENSION_H',),
                                    ('TORCH_EXTENSION_NAME',
                                     name.split('.')[-1]),
                                    ('_GLIBCXX_USE_CXX11_ABI',
                                     str(int(torch._C._GLIBCXX_USE_CXX11_ABI)))
                                ],
                                extra_compile_args=['-std=c++17', '-fopenmp'],
                                extra_link_args=['-std=c++17', '-fopenmp'],
                                include_dirs=include_dirs,
                                library_dirs=library_dirs,
                                runtime_library_dirs=library_dirs)
            cls.modules.append(ext)
            return ext
        except ImportError:
            return None

    @classmethod
    def Tensorflow(cls, name, sources):
        try:
            import tensorflow as tf
            ext = Extension(name,
                            sources=sources,
                            libraries=['cutensor', 'cudart'],
                            extra_compile_args=tf.sysconfig.get_compile_flags(),
                            extra_link_args=tf.sysconfig.get_link_flags() +
                            tf.sysconfig.get_compile_flags(),
                            define_macros=[('GOOGLE_CUDA', '1')],
                            include_dirs=include_dirs,
                            library_dirs=library_dirs,
                            runtime_library_dirs=library_dirs)
            cls.modules.append(ext)
        except ImportError:
            return None


Tensorflow

                                sources=sources,
                                libraries=['cutensor'],
                                define_macros=[
                                    ('TORCH_API_INCLUDE_EXTENSION_H',),
                                    ('TORCH_EXTENSION_NAME',
                                     name.split('.')[-1]),
                                    ('_GLIBCXX_USE_CXX11_ABI',
                                     str(int(torch._C._GLIBCXX_USE_CXX11_ABI)))
                                ],
                                extra_compile_args=['-std=c++17', '-fopenmp'],
                                extra_link_args=['-std=c++17', '-fopenmp'],
                                include_dirs=include_dirs,
                                library_dirs=library_dirs,
                                runtime_library_dirs=library_dirs)
            cls.modules.append(ext)
            return ext
        except ImportError:
            return None

    @classmethod
    def Tensorflow(cls, name, sources):
        try:
            import tensorflow as tf
            ext = Extension(name,
                            sources=sources,
                            libraries=['cutensor', 'cudart'],
                            extra_compile_args=tf.sysconfig.get_compile_flags(),
                            extra_link_args=tf.sysconfig.get_link_flags() +
                            tf.sysconfig.get_compile_flags(),
                            define_macros=[('GOOGLE_CUDA', '1')],
                            include_dirs=include_dirs,
                            library_dirs=library_dirs,
                            runtime_library_dirs=library_dirs)
            cls.modules.append(ext)
        except ImportError:
            return None


#! /usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name(s) of the copyright holder(s) nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

MAJOR = 0
MINOR = 1
PATCH = 0

VERSION = (MAJOR, MINOR, PATCH)

__version__ = '.'.join(map(str, VERSION))

__package_name__ = 'cutensor-python'
__description__ = 'PyTorch and Tensorflow Python bindings for cuTENSOR',
__homepage__ = 'https://developer.nvidia.com/cutensor',
__download_url__ = 'https://github.com/NVIDIA/CUDALibrarySamples/tree/master/cuTENSOR/cutensor',
__license__ = 'BSD'


#! /usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name(s) of the copyright holder(s) nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import torch

import unittest
from parameterized import parameterized
from parameterized import param

import cutensor.torch as cutensor

class EinsumTest(unittest.TestCase):

    def setUp(self):
        torch.backends.cuda.matmul.allow_tf32 = False

    def assertClose(self, cutensor_tensor, torch_tensor):
        self.assertEqual(cutensor_tensor.shape, torch_tensor.shape)
        self.assertEqual(torch.is_complex(cutensor_tensor), torch.is_complex(torch_tensor))
        if torch.is_complex(cutensor_tensor):
            self.assertClose(torch.real(cutensor_tensor), torch.real(torch_tensor))
            self.assertClose(torch.imag(cutensor_tensor), torch.imag(torch_tensor))
        else:
            torch.testing.assert_close(cutensor_tensor, torch_tensor, rtol=5e-3, atol=6e-3)

    @parameterized.expand(
        # yapf: disable
        [
            param(
                "test 0",
                a_size=(48, 37),
                b_size=(37, 74),
                equation="ik,kj->ij",
                dtype=torch.float32,
            ),
            param(
                "test 0 (complex)",
                a_size=(50, 50),
                b_size=(50, 50),
                equation="ik,kj->ij",
                dtype=torch.complex64,
            ),
            param(
                "test 1",
                a_size=(50, 50, 50),
                b_size=(50, 50, 50),
                equation="lik,lkj->lij",
                dtype=torch.complex128,
            ),
            param(
                "test 2",
                a_size=(50, 50, 50, 20),
                b_size=(50, 50, 50, 20),
                equation="likm,lkjm->lij",
                dtype=torch.float32,
            ),
            param(
                "test 3",
                a_size=(20, 50, 50, 50),
                b_size=(50, 50, 50, 20),
                equation="mlik,lkjm->lij",
                dtype=torch.float32,
            ),
            param(
                "test 4",
                a_size=(50, 50),
                b_size=(50, 50),
                equation="ik,kj->ij",
                dtype=torch.float16,
            ),
            param("test 5",
                  a_size=(50, 50, 50),
                  b_size=(50, 50, 50),
                  equation="lik,lkj->lij",
                  dtype=torch.float16),
            param(
                "test 6",
                a_size=(50, 50, 50, 20),
                b_size=(50, 50, 50, 20),
                equation="likm,lkjm->lij",
                dtype=torch.float16,
            ),
            param(
                "test 7",
                a_size=(20, 50, 50, 50),
                b_size=(50, 50, 50, 20),
                equation="mlik,lkjm->lij",
                dtype=torch.float16,
            ),
            param(
                "test 8",
                a_size=(2, 5, 50, 2),
                b_size=(5, 2, 50, 2),
                equation="mlik,lkjm",
                dtype=torch.float64,
            ),
            # Activate when cuTENSOR supports it
            # param(
            #     "test 8",
            #     a_size=(20, 50, 50, 50),
            #     b_size=(50, 50, 50, 20),
            #     equation="mlik,lkjm->lij",
            #     dtype=torch.bfloat16,
            # ),
        ]
        # yapf: enable
    )
    def test_einsum_equivalent_results(self,
                                       _,
                                       a_size,
                                       b_size,
                                       equation,
                                       dtype=torch.float32):

        kwargs = {
            'dtype': dtype,
            'device': torch.device("cuda"),
            'requires_grad': True
        }

        torch.manual_seed(0)

        cutensor_A = torch.randn(*a_size, **kwargs)
        cutensor_B = torch.randn(*b_size, **kwargs)
        cutensor_rslt = cutensor.EinsumFunction.apply(equation, cutensor_A,
                                                      cutensor_B)
        cutensor_rslt.backward(torch.ones_like(cutensor_rslt))
        cutensor_rslt = cutensor_rslt
        cutensor_A_grad = cutensor_A.grad
        cutensor_B_grad = cutensor_B.grad

        torch_A = cutensor_A.clone().detach().requires_grad_(True)
        torch_B = cutensor_B.clone().detach().requires_grad_(True)
        torch_rslt = torch.einsum(equation, torch_A, torch_B)
        torch_rslt.backward(torch.ones_like(torch_rslt))
        torch_A_grad = torch_A.grad
        torch_B_grad = torch_B.grad

        self.assertClose(cutensor_rslt, torch_rslt)
        self.assertClose(cutensor_A_grad, torch_A_grad)
        self.assertClose(cutensor_B_grad, torch_B_grad)

    @parameterized.expand(
        # yapf: disable
        [
            param(
                "test 0",
                sizes=[(50, 60), (60, 40)],
                equation="ik,kj->ji",
                dtype=torch.float32,
            ),
            param(
                "test 1",
                sizes=[(50, 60), (60, 7), (7, 8)],
                equation="ik,kl,lj->ij",
                dtype=torch.float32,
            ),
            param(
                "test 2",
                sizes=[(50, 60), (60, 7), (7, 8)],
                equation="ik,kl,lj",
                dtype=torch.float32,
            ),
            param(
                "test 3",
                sizes=[(50, 60), (60, 7), (7, 8)],
                equation="ik,kl,lj->ij",
                dtype=torch.complex64,
            ),
            # single input currently not supported
            param(
                "test 4",
                sizes=[(50, 60)],
                equation="ij->ji",
                dtype=torch.float32,
            ),
        ]
        # yapf: enable
    )
    def test_einsum_general_equivalent_results(self,
                                               _,
                                               sizes,
                                               equation,
                                               dtype=torch.float32):

        kwargs = {
            'dtype': dtype,
            'device': torch.device("cuda"),
            'requires_grad': True
        }

        cutensor_tensors = [torch.randn(*size, **kwargs) for size in sizes]
        torch_tensors = [
            t.clone().detach().requires_grad_(True) for t in cutensor_tensors
        ]

        cutensor_rslt = cutensor.EinsumGeneral(equation, *cutensor_tensors)
        cutensor_rslt.backward(torch.ones_like(cutensor_rslt))
        cutensor_rslt = cutensor_rslt
        cutensor_grads = [
            t.grad for t in cutensor_tensors
        ]

        torch_rslt = torch.einsum(equation, *torch_tensors)
        torch_rslt.backward(torch.ones_like(torch_rslt))
        torch_rslt = torch_rslt
        torch_grads = [t.grad for t in torch_tensors]

        self.assertClose(cutensor_rslt, torch_rslt)
        for ct, tt in zip(cutensor_grads, torch_grads):
            self.assertClose(ct, tt)

if __name__ == '__main__':
    unittest.main()


EinsumTest

# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import torch

import unittest
from parameterized import parameterized
from parameterized import param

import cutensor.torch as cutensor

class EinsumTest(unittest.TestCase):

    def setUp(self):
        torch.backends.cuda.matmul.allow_tf32 = False

    def assertClose(self, cutensor_tensor, torch_tensor):
        self.assertEqual(cutensor_tensor.shape, torch_tensor.shape)
        self.assertEqual(torch.is_complex(cutensor_tensor), torch.is_complex(torch_tensor))
        if torch.is_complex(cutensor_tensor):
            self.assertClose(torch.real(cutensor_tensor), torch.real(torch_tensor))
            self.assertClose(torch.imag(cutensor_tensor), torch.imag(torch_tensor))
        else:
            torch.testing.assert_close(cutensor_tensor, torch_tensor, rtol=5e-3, atol=6e-3)

    @parameterized.expand(
        # yapf: disable
        [
            param(
                "test 0",
                a_size=(48, 37),
                b_size=(37, 74),
                equation="ik,kj->ij",
                dtype=torch.float32,
            ),
            param(
                "test 0 (complex)",
                a_size=(50, 50),
                b_size=(50, 50),
                equation="ik,kj->ij",
                dtype=torch.complex64,
            ),
            param(
                "test 1",
                a_size=(50, 50, 50),
                b_size=(50, 50, 50),
                equation="lik,lkj->lij",
                dtype=torch.complex128,
            ),
            param(
                "test 2",
                a_size=(50, 50, 50, 20),
                b_size=(50, 50, 50, 20),
                equation="likm,lkjm->lij",
                dtype=torch.float32,
            ),
            param(
                "test 3",
                a_size=(20, 50, 50, 50),
                b_size=(50, 50, 50, 20),
                equation="mlik,lkjm->lij",
                dtype=torch.float32,
            ),
            param(
                "test 4",
                a_size=(50, 50),
                b_size=(50, 50),
                equation="ik,kj->ij",
                dtype=torch.float16,
            ),
            param("test 5",
                  a_size=(50, 50, 50),
                  b_size=(50, 50, 50),
                  equation="lik,lkj->lij",
                  dtype=torch.float16),
            param(
                "test 6",
                a_size=(50, 50, 50, 20),
                b_size=(50, 50, 50, 20),
                equation="likm,lkjm->lij",
                dtype=torch.float16,
            ),
            param(
                "test 7",
                a_size=(20, 50, 50, 50),
                b_size=(50, 50, 50, 20),
                equation="mlik,lkjm->lij",
                dtype=torch.float16,
            ),
            param(
                "test 8",
                a_size=(2, 5, 50, 2),
                b_size=(5, 2, 50, 2),
                equation="mlik,lkjm",
                dtype=torch.float64,
            ),
            # Activate when cuTENSOR supports it
            # param(
            #     "test 8",
            #     a_size=(20, 50, 50, 50),
            #     b_size=(50, 50, 50, 20),
            #     equation="mlik,lkjm->lij",
            #     dtype=torch.bfloat16,
            # ),
        ]
        # yapf: enable
    )
    def test_einsum_equivalent_results(self,
                                       _,
                                       a_size,
                                       b_size,
                                       equation,
                                       dtype=torch.float32):

        kwargs = {
            'dtype': dtype,
            'device': torch.device("cuda"),
            'requires_grad': True
        }

        torch.manual_seed(0)

        cutensor_A = torch.randn(*a_size, **kwargs)
        cutensor_B = torch.randn(*b_size, **kwargs)
        cutensor_rslt = cutensor.EinsumFunction.apply(equation, cutensor_A,
                                                      cutensor_B)
        cutensor_rslt.backward(torch.ones_like(cutensor_rslt))
        cutensor_rslt = cutensor_rslt
        cutensor_A_grad = cutensor_A.grad
        cutensor_B_grad = cutensor_B.grad

        torch_A = cutensor_A.clone().detach().requires_grad_(True)
        torch_B = cutensor_B.clone().detach().requires_grad_(True)
        torch_rslt = torch.einsum(equation, torch_A, torch_B)
        torch_rslt.backward(torch.ones_like(torch_rslt))
        torch_A_grad = torch_A.grad
        torch_B_grad = torch_B.grad

        self.assertClose(cutensor_rslt, torch_rslt)
        self.assertClose(cutensor_A_grad, torch_A_grad)
        self.assertClose(cutensor_B_grad, torch_B_grad)

    @parameterized.expand(
        # yapf: disable
        [
            param(
                "test 0",
                sizes=[(50, 60), (60, 40)],
                equation="ik,kj->ji",
                dtype=torch.float32,
            ),
            param(
                "test 1",
                sizes=[(50, 60), (60, 7), (7, 8)],
                equation="ik,kl,lj->ij",
                dtype=torch.float32,
            ),
            param(
                "test 2",
                sizes=[(50, 60), (60, 7), (7, 8)],
                equation="ik,kl,lj",
                dtype=torch.float32,
            ),
            param(
                "test 3",
                sizes=[(50, 60), (60, 7), (7, 8)],
                equation="ik,kl,lj->ij",
                dtype=torch.complex64,
            ),
            # single input currently not supported
            param(
                "test 4",
                sizes=[(50, 60)],
                equation="ij->ji",
                dtype=torch.float32,
            ),
        ]
        # yapf: enable
    )
    def test_einsum_general_equivalent_results(self,
                                               _,
                                               sizes,
                                               equation,
                                               dtype=torch.float32):

        kwargs = {
            'dtype': dtype,
            'device': torch.device("cuda"),
            'requires_grad': True
        }

        cutensor_tensors = [torch.randn(*size, **kwargs) for size in sizes]
        torch_tensors = [
            t.clone().detach().requires_grad_(True) for t in cutensor_tensors
        ]

        cutensor_rslt = cutensor.EinsumGeneral(equation, *cutensor_tensors)
        cutensor_rslt.backward(torch.ones_like(cutensor_rslt))
        cutensor_rslt = cutensor_rslt
        cutensor_grads = [
            t.grad for t in cutensor_tensors
        ]

        torch_rslt = torch.einsum(equation, *torch_tensors)
        torch_rslt.backward(torch.ones_like(torch_rslt))
        torch_rslt = torch_rslt
        torch_grads = [t.grad for t in torch_tensors]

        self.assertClose(cutensor_rslt, torch_rslt)
        for ct, tt in zip(cutensor_grads, torch_grads):
            self.assertClose(ct, tt)

if __name__ == '__main__':
    unittest.main()


setUp

# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import torch

import unittest
from parameterized import parameterized
from parameterized import param

import cutensor.torch as cutensor

class EinsumTest(unittest.TestCase):

    def setUp(self):
        torch.backends.cuda.matmul.allow_tf32 = False

    def assertClose(self, cutensor_tensor, torch_tensor):
        self.assertEqual(cutensor_tensor.shape, torch_tensor.shape)
        self.assertEqual(torch.is_complex(cutensor_tensor), torch.is_complex(torch_tensor))
        if torch.is_complex(cutensor_tensor):
            self.assertClose(torch.real(cutensor_tensor), torch.real(torch_tensor))
            self.assertClose(torch.imag(cutensor_tensor), torch.imag(torch_tensor))
        else:
            torch.testing.assert_close(cutensor_tensor, torch_tensor, rtol=5e-3, atol=6e-3)

    @parameterized.expand(
        # yapf: disable
        [
            param(
                "test 0",
                a_size=(48, 37),
                b_size=(37, 74),
                equation="ik,kj->ij",


assertClose

# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import torch

import unittest
from parameterized import parameterized
from parameterized import param

import cutensor.torch as cutensor

class EinsumTest(unittest.TestCase):

    def setUp(self):
        torch.backends.cuda.matmul.allow_tf32 = False

    def assertClose(self, cutensor_tensor, torch_tensor):
        self.assertEqual(cutensor_tensor.shape, torch_tensor.shape)
        self.assertEqual(torch.is_complex(cutensor_tensor), torch.is_complex(torch_tensor))
        if torch.is_complex(cutensor_tensor):
            self.assertClose(torch.real(cutensor_tensor), torch.real(torch_tensor))
            self.assertClose(torch.imag(cutensor_tensor), torch.imag(torch_tensor))
        else:
            torch.testing.assert_close(cutensor_tensor, torch_tensor, rtol=5e-3, atol=6e-3)

    @parameterized.expand(
        # yapf: disable
        [
            param(
                "test 0",
                a_size=(48, 37),
                b_size=(37, 74),
                equation="ik,kj->ij",
                dtype=torch.float32,
            ),
            param(
                "test 0 (complex)",
                a_size=(50, 50),
                b_size=(50, 50),
                equation="ik,kj->ij",
                dtype=torch.complex64,
            ),
            param(


                dtype=torch.float16,
            ),
            param(
                "test 8",
                a_size=(2, 5, 50, 2),
                b_size=(5, 2, 50, 2),
                equation="mlik,lkjm",
                dtype=torch.float64,
            ),
            # Activate when cuTENSOR supports it
            # param(
            #     "test 8",
            #     a_size=(20, 50, 50, 50),
            #     b_size=(50, 50, 50, 20),
            #     equation="mlik,lkjm->lij",
            #     dtype=torch.bfloat16,
            # ),
        ]
        # yapf: enable
    )
    def test_einsum_equivalent_results(self,
                                       _,
                                       a_size,
                                       b_size,
                                       equation,
                                       dtype=torch.float32):

        kwargs = {
            'dtype': dtype,
            'device': torch.device("cuda"),
            'requires_grad': True
        }

        torch.manual_seed(0)

        cutensor_A = torch.randn(*a_size, **kwargs)
        cutensor_B = torch.randn(*b_size, **kwargs)
        cutensor_rslt = cutensor.EinsumFunction.apply(equation, cutensor_A,
                                                      cutensor_B)
        cutensor_rslt.backward(torch.ones_like(cutensor_rslt))
        cutensor_rslt = cutensor_rslt
        cutensor_A_grad = cutensor_A.grad
        cutensor_B_grad = cutensor_B.grad

        torch_A = cutensor_A.clone().detach().requires_grad_(True)
        torch_B = cutensor_B.clone().detach().requires_grad_(True)
        torch_rslt = torch.einsum(equation, torch_A, torch_B)
        torch_rslt.backward(torch.ones_like(torch_rslt))
        torch_A_grad = torch_A.grad
        torch_B_grad = torch_B.grad

        self.assertClose(cutensor_rslt, torch_rslt)
        self.assertClose(cutensor_A_grad, torch_A_grad)
        self.assertClose(cutensor_B_grad, torch_B_grad)

    @parameterized.expand(
        # yapf: disable
        [
            param(
                "test 0",
                sizes=[(50, 60), (60, 40)],
                equation="ik,kj->ji",
                dtype=torch.float32,
            ),
            param(
                "test 1",
                sizes=[(50, 60), (60, 7), (7, 8)],
                equation="ik,kl,lj->ij",
                dtype=torch.float32,
            ),
            param(
                "test 2",
                sizes=[(50, 60), (60, 7), (7, 8)],
                equation="ik,kl,lj",


                sizes=[(50, 60), (60, 7), (7, 8)],
                equation="ik,kl,lj",
                dtype=torch.float32,
            ),
            param(
                "test 3",
                sizes=[(50, 60), (60, 7), (7, 8)],
                equation="ik,kl,lj->ij",
                dtype=torch.complex64,
            ),
            # single input currently not supported
            param(
                "test 4",
                sizes=[(50, 60)],
                equation="ij->ji",
                dtype=torch.float32,
            ),
        ]
        # yapf: enable
    )
    def test_einsum_general_equivalent_results(self,
                                               _,
                                               sizes,
                                               equation,
                                               dtype=torch.float32):

        kwargs = {
            'dtype': dtype,
            'device': torch.device("cuda"),
            'requires_grad': True
        }

        cutensor_tensors = [torch.randn(*size, **kwargs) for size in sizes]
        torch_tensors = [
            t.clone().detach().requires_grad_(True) for t in cutensor_tensors
        ]

        cutensor_rslt = cutensor.EinsumGeneral(equation, *cutensor_tensors)
        cutensor_rslt.backward(torch.ones_like(cutensor_rslt))
        cutensor_rslt = cutensor_rslt
        cutensor_grads = [
            t.grad for t in cutensor_tensors
        ]

        torch_rslt = torch.einsum(equation, *torch_tensors)
        torch_rslt.backward(torch.ones_like(torch_rslt))
        torch_rslt = torch_rslt
        torch_grads = [t.grad for t in torch_tensors]

        self.assertClose(cutensor_rslt, torch_rslt)
        for ct, tt in zip(cutensor_grads, torch_grads):
            self.assertClose(ct, tt)

if __name__ == '__main__':
    unittest.main()


cuTENSOR.python.cutensor.torch.einsum

#! /usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name(s) of the copyright holder(s) nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import torch
import torch.autograd
import numpy as np
from .binding import einsum, plan, execute
from ..common import normalize_subscript

class EinsumFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, equation, input_0, input_1=None):
        equation, isBinary = normalize_subscript(equation)
        if isBinary and input_1 is None:
            raise RuntimeError('The subscript indicates two inputs, but only one was passed')
        if not isBinary and input_1 is not None:
            raise RuntimeError('The subscript indicates one input, but two were passed')
        if input_1 is None:
            input_1 = input_0.new_empty((1,))

        output = einsum(equation, input_0, input_1, False, False)

        if isBinary:
            ctx.save_for_backward(input_0, input_1)

        ctx.equation = equation
        ctx.isBinary = isBinary

        return output

    @staticmethod
    def backward(ctx, grad_output):
        equation = ctx.equation
        lhs, modeC = equation.split('->')
        if ctx.isBinary:
            input_0, input_1 = ctx.saved_tensors
            conjugate = False
            if torch.is_complex(input_0) or torch.is_complex(input_1):
                conjugate = True
            modeA, modeB = lhs.split(',')
            d_input_0 = einsum(modeC + ',' + modeB + '->' + modeA, grad_output,
                               input_1, False, conjugate)
            d_input_1 = einsum(modeA + ',' + modeC + '->' + modeB, input_0,
                               grad_output, conjugate, False)
            return None, d_input_0, d_input_1
        else:
            dummy = grad_output.new_empty((1,))
            d_input = einsum(modeC + '->' + lhs, grad_output, dummy, False, False)
            return None, d_input

    def plan(equation, input_0, input_1=None, jit_pref=False):
        equation, isBinary = normalize_subscript(equation)
        if isBinary and input_1 is None:
            raise RuntimeError('The subscript indicates two inputs, but only one was passed')
        if not isBinary and input_1 is not None:
            raise RuntimeError('The subscript indicates one input, but two were passed')
        if input_1 is None:
            input_1 = input_0.new_empty((1,))

        output = plan(equation, input_0, input_1, False, False, jit_pref)

        return output

    def execute(plan):
        try:
            result = execute(plan)
            if result is None:  # Handle NULL return
                raise RuntimeError("cuTENSOR execute returned NULL (CUTENSOR_STATUS_INVALID_VALUE)")
            return result
        except SystemError as e:
            raise RuntimeError(f"cuTENSOR execution failed {e}")

class Einsum(torch.nn.Module):

    def __init__(self, equation):
        super(Einsum, self).__init__()
        self.equation = equation
        self.reset_parameters()

    def reset_parameters(self):
        pass

    def forward(self, input_0, input_1):
        return EinsumFunction.apply(self.equation, input_0, input_1)

    def plan(self, input_0, input_1, jit_pref=False):
        return EinsumFunction.plan(self.equation, input_0, input_1=input_1, jit_pref=jit_pref)

    def execute(self, plan):
        return EinsumFunction.execute(plan)

def _compute_target_tensor(in0, in1, target, rest):
    result = ""
    rest = ''.join(rest) + target
    for m in in0[:-1] + in1[:-1] + in1[-1] + in0[-1]:
        if m in rest and not m in result:
            result += m
    # reorder target modes like target
    result = list(result)
    for i in range(len(result)):
        if result[i] not in target: continue
        for j in range(i):
            if result[j] not in target: continue
            if target.index(result[j]) > target.index(result[i]):
                result[i], result[j] = result[j], result[i]
    return ''.join(result)

def EinsumGeneral(equation, *tensors, **kwargs):
    tensors = list(tensors)
    equation, isBinary = normalize_subscript(equation)
    path = np.einsum_path(equation,
                          *[np.broadcast_to(np.nan, t.shape) for t in tensors],
                          **kwargs)
    path = path[0][1:]
    equation = equation.split('->')
    eqs = equation[0].split(',')
    target = equation[1]
    for step in path:
        if len(step) == 1:
            result = EinsumFunction.apply(eqs[0] + '->' + target, tensors[0])
            continue
        assert step[0] < step[1]
        in0 = tensors[step[0]]
        in1 = tensors[step[1]]
        tensors.pop(step[1])
        tensors.pop(step[0])
        tgt = _compute_target_tensor(eqs[step[0]], eqs[step[1]], target, [eq for idx, eq in enumerate(eqs) if idx not in step])
        assert tgt != ""
        eq = eqs[step[0]] + ',' + eqs[step[1]] + '->' + tgt
        eqs.pop(step[1])
        eqs.pop(step[0])
        eqs.append(tgt)
        result = EinsumFunction.apply(eq, in0, in1)
        tensors.append(result)
    return result


EinsumFunction

#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import torch
import torch.autograd
import numpy as np
from .binding import einsum, plan, execute
from ..common import normalize_subscript

class EinsumFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, equation, input_0, input_1=None):
        equation, isBinary = normalize_subscript(equation)
        if isBinary and input_1 is None:
            raise RuntimeError('The subscript indicates two inputs, but only one was passed')
        if not isBinary and input_1 is not None:
            raise RuntimeError('The subscript indicates one input, but two were passed')
        if input_1 is None:
            input_1 = input_0.new_empty((1,))

        output = einsum(equation, input_0, input_1, False, False)

        if isBinary:
            ctx.save_for_backward(input_0, input_1)

        ctx.equation = equation
        ctx.isBinary = isBinary

        return output

    @staticmethod
    def backward(ctx, grad_output):
        equation = ctx.equation
        lhs, modeC = equation.split('->')
        if ctx.isBinary:
            input_0, input_1 = ctx.saved_tensors
            conjugate = False
            if torch.is_complex(input_0) or torch.is_complex(input_1):
                conjugate = True
            modeA, modeB = lhs.split(',')
            d_input_0 = einsum(modeC + ',' + modeB + '->' + modeA, grad_output,
                               input_1, False, conjugate)
            d_input_1 = einsum(modeA + ',' + modeC + '->' + modeB, input_0,
                               grad_output, conjugate, False)
            return None, d_input_0, d_input_1
        else:
            dummy = grad_output.new_empty((1,))
            d_input = einsum(modeC + '->' + lhs, grad_output, dummy, False, False)
            return None, d_input

    def plan(equation, input_0, input_1=None, jit_pref=False):
        equation, isBinary = normalize_subscript(equation)
        if isBinary and input_1 is None:
            raise RuntimeError('The subscript indicates two inputs, but only one was passed')
        if not isBinary and input_1 is not None:
            raise RuntimeError('The subscript indicates one input, but two were passed')
        if input_1 is None:
            input_1 = input_0.new_empty((1,))

        output = plan(equation, input_0, input_1, False, False, jit_pref)

        return output

    def execute(plan):
        try:
            result = execute(plan)
            if result is None:  # Handle NULL return
                raise RuntimeError("cuTENSOR execute returned NULL (CUTENSOR_STATUS_INVALID_VALUE)")
            return result
        except SystemError as e:
            raise RuntimeError(f"cuTENSOR execution failed {e}")

class Einsum(torch.nn.Module):

    def __init__(self, equation):
        super(Einsum, self).__init__()
        self.equation = equation
        self.reset_parameters()

    def reset_parameters(self):
        pass

    def forward(self, input_0, input_1):
        return EinsumFunction.apply(self.equation, input_0, input_1)

    def plan(self, input_0, input_1, jit_pref=False):
        return EinsumFunction.plan(self.equation, input_0, input_1=input_1, jit_pref=jit_pref)

    def execute(self, plan):
        return EinsumFunction.execute(plan)


Einsum

        if not isBinary and input_1 is not None:
            raise RuntimeError('The subscript indicates one input, but two were passed')
        if input_1 is None:
            input_1 = input_0.new_empty((1,))

        output = plan(equation, input_0, input_1, False, False, jit_pref)

        return output

    def execute(plan):
        try:
            result = execute(plan)
            if result is None:  # Handle NULL return
                raise RuntimeError("cuTENSOR execute returned NULL (CUTENSOR_STATUS_INVALID_VALUE)")
            return result
        except SystemError as e:
            raise RuntimeError(f"cuTENSOR execution failed {e}")

class Einsum(torch.nn.Module):

    def __init__(self, equation):
        super(Einsum, self).__init__()
        self.equation = equation
        self.reset_parameters()

    def reset_parameters(self):
        pass

    def forward(self, input_0, input_1):
        return EinsumFunction.apply(self.equation, input_0, input_1)

    def plan(self, input_0, input_1, jit_pref=False):
        return EinsumFunction.plan(self.equation, input_0, input_1=input_1, jit_pref=jit_pref)

    def execute(self, plan):
        return EinsumFunction.execute(plan)

def _compute_target_tensor(in0, in1, target, rest):
    result = ""
    rest = ''.join(rest) + target
    for m in in0[:-1] + in1[:-1] + in1[-1] + in0[-1]:
        if m in rest and not m in result:
            result += m
    # reorder target modes like target
    result = list(result)
    for i in range(len(result)):
        if result[i] not in target: continue
        for j in range(i):
            if result[j] not in target: continue
            if target.index(result[j]) > target.index(result[i]):
                result[i], result[j] = result[j], result[i]
    return ''.join(result)

def EinsumGeneral(equation, *tensors, **kwargs):


class Einsum(torch.nn.Module):

    def __init__(self, equation):
        super(Einsum, self).__init__()
        self.equation = equation
        self.reset_parameters()

    def reset_parameters(self):
        pass

    def forward(self, input_0, input_1):
        return EinsumFunction.apply(self.equation, input_0, input_1)

    def plan(self, input_0, input_1, jit_pref=False):
        return EinsumFunction.plan(self.equation, input_0, input_1=input_1, jit_pref=jit_pref)

    def execute(self, plan):
        return EinsumFunction.execute(plan)

def _compute_target_tensor(in0, in1, target, rest):
    result = ""
    rest = ''.join(rest) + target
    for m in in0[:-1] + in1[:-1] + in1[-1] + in0[-1]:
        if m in rest and not m in result:
            result += m
    # reorder target modes like target
    result = list(result)
    for i in range(len(result)):
        if result[i] not in target: continue
        for j in range(i):
            if result[j] not in target: continue
            if target.index(result[j]) > target.index(result[i]):
                result[i], result[j] = result[j], result[i]
    return ''.join(result)

def EinsumGeneral(equation, *tensors, **kwargs):
    tensors = list(tensors)
    equation, isBinary = normalize_subscript(equation)
    path = np.einsum_path(equation,
                          *[np.broadcast_to(np.nan, t.shape) for t in tensors],
                          **kwargs)
    path = path[0][1:]
    equation = equation.split('->')
    eqs = equation[0].split(',')
    target = equation[1]
    for step in path:
        if len(step) == 1:
            result = EinsumFunction.apply(eqs[0] + '->' + target, tensors[0])
            continue
        assert step[0] < step[1]
        in0 = tensors[step[0]]
        in1 = tensors[step[1]]
        tensors.pop(step[1])


EinsumGeneral

        return EinsumFunction.execute(plan)

def _compute_target_tensor(in0, in1, target, rest):
    result = ""
    rest = ''.join(rest) + target
    for m in in0[:-1] + in1[:-1] + in1[-1] + in0[-1]:
        if m in rest and not m in result:
            result += m
    # reorder target modes like target
    result = list(result)
    for i in range(len(result)):
        if result[i] not in target: continue
        for j in range(i):
            if result[j] not in target: continue
            if target.index(result[j]) > target.index(result[i]):
                result[i], result[j] = result[j], result[i]
    return ''.join(result)

def EinsumGeneral(equation, *tensors, **kwargs):
    tensors = list(tensors)
    equation, isBinary = normalize_subscript(equation)
    path = np.einsum_path(equation,
                          *[np.broadcast_to(np.nan, t.shape) for t in tensors],
                          **kwargs)
    path = path[0][1:]
    equation = equation.split('->')
    eqs = equation[0].split(',')
    target = equation[1]
    for step in path:
        if len(step) == 1:
            result = EinsumFunction.apply(eqs[0] + '->' + target, tensors[0])
            continue
        assert step[0] < step[1]
        in0 = tensors[step[0]]
        in1 = tensors[step[1]]
        tensors.pop(step[1])
        tensors.pop(step[0])
        tgt = _compute_target_tensor(eqs[step[0]], eqs[step[1]], target, [eq for idx, eq in enumerate(eqs) if idx not in step])
        assert tgt != ""
        eq = eqs[step[0]] + ',' + eqs[step[1]] + '->' + tgt
        eqs.pop(step[1])
        eqs.pop(step[0])
        eqs.append(tgt)
        result = EinsumFunction.apply(eq, in0, in1)
        tensors.append(result)
    return result


forward

    def execute(plan):
        try:
            result = execute(plan)
            if result is None:  # Handle NULL return
                raise RuntimeError("cuTENSOR execute returned NULL (CUTENSOR_STATUS_INVALID_VALUE)")
            return result
        except SystemError as e:
            raise RuntimeError(f"cuTENSOR execution failed {e}")

class Einsum(torch.nn.Module):

    def __init__(self, equation):
        super(Einsum, self).__init__()
        self.equation = equation
        self.reset_parameters()

    def reset_parameters(self):
        pass

    def forward(self, input_0, input_1):
        return EinsumFunction.apply(self.equation, input_0, input_1)

    def plan(self, input_0, input_1, jit_pref=False):
        return EinsumFunction.plan(self.equation, input_0, input_1=input_1, jit_pref=jit_pref)

    def execute(self, plan):
        return EinsumFunction.execute(plan)

def _compute_target_tensor(in0, in1, target, rest):
    result = ""
    rest = ''.join(rest) + target
    for m in in0[:-1] + in1[:-1] + in1[-1] + in0[-1]:
        if m in rest and not m in result:
            result += m
    # reorder target modes like target
    result = list(result)
    for i in range(len(result)):
        if result[i] not in target: continue
        for j in range(i):
            if result[j] not in target: continue


backward

    def forward(ctx, equation, input_0, input_1=None):
        equation, isBinary = normalize_subscript(equation)
        if isBinary and input_1 is None:
            raise RuntimeError('The subscript indicates two inputs, but only one was passed')
        if not isBinary and input_1 is not None:
            raise RuntimeError('The subscript indicates one input, but two were passed')
        if input_1 is None:
            input_1 = input_0.new_empty((1,))

        output = einsum(equation, input_0, input_1, False, False)

        if isBinary:
            ctx.save_for_backward(input_0, input_1)

        ctx.equation = equation
        ctx.isBinary = isBinary

        return output

    @staticmethod
    def backward(ctx, grad_output):
        equation = ctx.equation
        lhs, modeC = equation.split('->')
        if ctx.isBinary:
            input_0, input_1 = ctx.saved_tensors
            conjugate = False
            if torch.is_complex(input_0) or torch.is_complex(input_1):
                conjugate = True
            modeA, modeB = lhs.split(',')
            d_input_0 = einsum(modeC + ',' + modeB + '->' + modeA, grad_output,
                               input_1, False, conjugate)
            d_input_1 = einsum(modeA + ',' + modeC + '->' + modeB, input_0,
                               grad_output, conjugate, False)
            return None, d_input_0, d_input_1
        else:
            dummy = grad_output.new_empty((1,))
            d_input = einsum(modeC + '->' + lhs, grad_output, dummy, False, False)
            return None, d_input

    def plan(equation, input_0, input_1=None, jit_pref=False):
        equation, isBinary = normalize_subscript(equation)
        if isBinary and input_1 is None:
            raise RuntimeError('The subscript indicates two inputs, but only one was passed')
        if not isBinary and input_1 is not None:
            raise RuntimeError('The subscript indicates one input, but two were passed')
        if input_1 is None:
            input_1 = input_0.new_empty((1,))

        output = plan(equation, input_0, input_1, False, False, jit_pref)

        return output

    def execute(plan):
        try:
            result = execute(plan)
            if result is None:  # Handle NULL return


plan

            if result is None:  # Handle NULL return
                raise RuntimeError("cuTENSOR execute returned NULL (CUTENSOR_STATUS_INVALID_VALUE)")
            return result
        except SystemError as e:
            raise RuntimeError(f"cuTENSOR execution failed {e}")

class Einsum(torch.nn.Module):

    def __init__(self, equation):
        super(Einsum, self).__init__()
        self.equation = equation
        self.reset_parameters()

    def reset_parameters(self):
        pass

    def forward(self, input_0, input_1):
        return EinsumFunction.apply(self.equation, input_0, input_1)

    def plan(self, input_0, input_1, jit_pref=False):
        return EinsumFunction.plan(self.equation, input_0, input_1=input_1, jit_pref=jit_pref)

    def execute(self, plan):
        return EinsumFunction.execute(plan)

def _compute_target_tensor(in0, in1, target, rest):
    result = ""
    rest = ''.join(rest) + target
    for m in in0[:-1] + in1[:-1] + in1[-1] + in0[-1]:
        if m in rest and not m in result:
            result += m
    # reorder target modes like target
    result = list(result)
    for i in range(len(result)):
        if result[i] not in target: continue
        for j in range(i):
            if result[j] not in target: continue
            if target.index(result[j]) > target.index(result[i]):
                result[i], result[j] = result[j], result[i]
    return ''.join(result)


execute

        except SystemError as e:
            raise RuntimeError(f"cuTENSOR execution failed {e}")

class Einsum(torch.nn.Module):

    def __init__(self, equation):
        super(Einsum, self).__init__()
        self.equation = equation
        self.reset_parameters()

    def reset_parameters(self):
        pass

    def forward(self, input_0, input_1):
        return EinsumFunction.apply(self.equation, input_0, input_1)

    def plan(self, input_0, input_1, jit_pref=False):
        return EinsumFunction.plan(self.equation, input_0, input_1=input_1, jit_pref=jit_pref)

    def execute(self, plan):
        return EinsumFunction.execute(plan)

def _compute_target_tensor(in0, in1, target, rest):
    result = ""
    rest = ''.join(rest) + target
    for m in in0[:-1] + in1[:-1] + in1[-1] + in0[-1]:
        if m in rest and not m in result:
            result += m
    # reorder target modes like target
    result = list(result)
    for i in range(len(result)):
        if result[i] not in target: continue
        for j in range(i):
            if result[j] not in target: continue
            if target.index(result[j]) > target.index(result[i]):
                result[i], result[j] = result[j], result[i]
    return ''.join(result)

def EinsumGeneral(equation, *tensors, **kwargs):


        if input_1 is None:
            input_1 = input_0.new_empty((1,))

        output = plan(equation, input_0, input_1, False, False, jit_pref)

        return output

    def execute(plan):
        try:
            result = execute(plan)
            if result is None:  # Handle NULL return
                raise RuntimeError("cuTENSOR execute returned NULL (CUTENSOR_STATUS_INVALID_VALUE)")
            return result
        except SystemError as e:
            raise RuntimeError(f"cuTENSOR execution failed {e}")

class Einsum(torch.nn.Module):

    def __init__(self, equation):
        super(Einsum, self).__init__()
        self.equation = equation
        self.reset_parameters()

    def reset_parameters(self):
        pass

    def forward(self, input_0, input_1):
        return EinsumFunction.apply(self.equation, input_0, input_1)

    def plan(self, input_0, input_1, jit_pref=False):
        return EinsumFunction.plan(self.equation, input_0, input_1=input_1, jit_pref=jit_pref)

    def execute(self, plan):
        return EinsumFunction.execute(plan)

def _compute_target_tensor(in0, in1, target, rest):
    result = ""
    rest = ''.join(rest) + target
    for m in in0[:-1] + in1[:-1] + in1[-1] + in0[-1]:
        if m in rest and not m in result:
            result += m


        return output

    def execute(plan):
        try:
            result = execute(plan)
            if result is None:  # Handle NULL return
                raise RuntimeError("cuTENSOR execute returned NULL (CUTENSOR_STATUS_INVALID_VALUE)")
            return result
        except SystemError as e:
            raise RuntimeError(f"cuTENSOR execution failed {e}")

class Einsum(torch.nn.Module):

    def __init__(self, equation):
        super(Einsum, self).__init__()
        self.equation = equation
        self.reset_parameters()

    def reset_parameters(self):
        pass

    def forward(self, input_0, input_1):
        return EinsumFunction.apply(self.equation, input_0, input_1)

    def plan(self, input_0, input_1, jit_pref=False):
        return EinsumFunction.plan(self.equation, input_0, input_1=input_1, jit_pref=jit_pref)

    def execute(self, plan):
        return EinsumFunction.execute(plan)

def _compute_target_tensor(in0, in1, target, rest):
    result = ""
    rest = ''.join(rest) + target
    for m in in0[:-1] + in1[:-1] + in1[-1] + in0[-1]:
        if m in rest and not m in result:
            result += m
    # reorder target modes like target
    result = list(result)
    for i in range(len(result)):


# ! /usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name(s) of the copyright holder(s) nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from parameterized import parameterized
from parameterized import param

import tensorflow as tf
from tensorflow.python.platform import test
import tensorflow.test

import cutensor.tensorflow as cutensor

tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

class EinsumcuTENSORTest(tensorflow.test.TestCase):

    def setUp(self):
        super().setUp()
        tf.config.experimental.enable_tensor_float_32_execution(False)

    @parameterized.expand(
        # yapf: disable
        [
            param(
                "test 0",
                a_size=(50, 50),
                b_size=(50, 50),
                equation="ik,kj->ij",
                dtype=tf.float32,
            ),
            param(
                "test 1",
                a_size=(50, 50, 50),
                b_size=(50, 50, 50),
                equation="lik,lkj->lij",
                dtype=tf.float32,
            ),
            param(
                "test 2",
                a_size=(50, 50, 50, 20),
                b_size=(50, 50, 50, 20),
                equation="likm,lkjm->lij",
                dtype=tf.float32,
            ),
            param(
                "test 3",
                a_size=(20, 50, 50, 50),
                b_size=(50, 50, 50, 20),
                equation="mlik,lkjm->lij",
                dtype=tf.float32,
            ),
            param(
                "test 4",
                a_size=(50, 50),
                b_size=(50, 50),
                equation="ik,kj->ij",
                dtype=tf.float16,
            ),
            param("test 5",
                  a_size=(50, 50, 50),
                  b_size=(50, 50, 50),
                  equation="lik,lkj->lij",
                  dtype=tf.float16),
            param(
                "test 6",
                a_size=(50, 50, 50, 20),
                b_size=(50, 50, 50, 20),
                equation="likm,lkjm->lij",
                dtype=tf.float16,
            ),
            param(
                "test 7",
                a_size=(20, 50, 50, 50),
                b_size=(50, 50, 50, 20),
                equation="mlik,lkjm->lij",
                dtype=tf.float16,
            ),
            param(
                "test 8",
                a_size=(2, 5, 5, 5),
                b_size=(5, 5, 5, 2),
                equation="mlik,lkjm",
                dtype=tf.float16,
            ),
            param(
                "test 9",
                a_size=(20, 50, 50, 50),
                b_size=None,
                equation="mlik->imlk",
                dtype=tf.float16,
            ),
            # Activate when cuTENSOR supports it
            # param(
            #     "test 8",
            #     a_size=(20, 50, 50, 50),
            #     b_size=(50, 50, 50, 20),
            #     equation="mlik,lkjm->lij",
            #     dtype=tf.bfloat16,
            # ),
        ]
        # yapf: enable
    )
    def test_einsum_equivalent_results(self,
                                       _,
                                       a_size,
                                       b_size,
                                       equation,
                                       dtype=tf.float32):
        A = tf.random.normal(a_size, dtype=dtype)

        if b_size is not None:
            B = tf.random.normal(b_size, dtype=dtype)

            with tf.GradientTape() as tape:
                tape.watch([A, B])
                tf_native_rslt = tf.einsum(equation, A, B, name="tf_native_einsum")

            tf_native_grads = tape.gradient(tf_native_rslt, [A, B])

            with tf.GradientTape() as tape:
                tape.watch([A, B])
                tf_cutensor_rslt = cutensor.einsum(equation,
                                                   A,
                                                   B,
                                                   name="tf_cuTensor_einsum")

            tf_cutensor_grads = tape.gradient(tf_cutensor_rslt, [A, B])
        else:
            with tf.GradientTape() as tape:
                tape.watch([A])
                tf_native_rslt = tf.einsum(equation, A, name="tf_native_einsum")

            tf_native_grads = tape.gradient(tf_native_rslt, [A])

            with tf.GradientTape() as tape:
                tape.watch([A])
                tf_cutensor_rslt = cutensor.einsum(equation,
                                                   A,
                                                   name="tf_cuTensor_einsum")
            tf_cutensor_grads = tape.gradient(tf_cutensor_rslt, [A])

        self.assertEqual(tf_native_rslt.get_shape(),
                         tf_cutensor_rslt.get_shape())

        self.assertEqual(tf_native_rslt.dtype, tf_cutensor_rslt.dtype)
        self.assertEqual(len(tf_cutensor_grads), len(tf_native_grads))

        self.assertAllClose(tf_native_rslt,
                            tf_cutensor_rslt,
                            rtol=5e-03,
                            atol=5e-03)

        for tf_native_grad, tf_cutensor_grad in zip(tf_native_grads,
                                                    tf_cutensor_grads):
            self.assertAllClose(tf_native_grad,
                                tf_cutensor_grad,
                                rtol=5e-03,
                                atol=5e-03)
            self.assertEqual(tf_native_grad.dtype, tf_cutensor_grad.dtype)

if __name__ == '__main__':
    test.main()


EinsumcuTENSORTest

# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from parameterized import parameterized
from parameterized import param

import tensorflow as tf
from tensorflow.python.platform import test
import tensorflow.test

import cutensor.tensorflow as cutensor

tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

class EinsumcuTENSORTest(tensorflow.test.TestCase):

    def setUp(self):
        super().setUp()
        tf.config.experimental.enable_tensor_float_32_execution(False)

    @parameterized.expand(
        # yapf: disable
        [
            param(
                "test 0",
                a_size=(50, 50),
                b_size=(50, 50),
                equation="ik,kj->ij",
                dtype=tf.float32,
            ),
            param(
                "test 1",
                a_size=(50, 50, 50),
                b_size=(50, 50, 50),
                equation="lik,lkj->lij",
                dtype=tf.float32,
            ),
            param(
                "test 2",
                a_size=(50, 50, 50, 20),
                b_size=(50, 50, 50, 20),
                equation="likm,lkjm->lij",
                dtype=tf.float32,
            ),
            param(
                "test 3",
                a_size=(20, 50, 50, 50),
                b_size=(50, 50, 50, 20),
                equation="mlik,lkjm->lij",
                dtype=tf.float32,
            ),
            param(
                "test 4",
                a_size=(50, 50),
                b_size=(50, 50),
                equation="ik,kj->ij",
                dtype=tf.float16,
            ),
            param("test 5",
                  a_size=(50, 50, 50),
                  b_size=(50, 50, 50),
                  equation="lik,lkj->lij",
                  dtype=tf.float16),
            param(
                "test 6",
                a_size=(50, 50, 50, 20),
                b_size=(50, 50, 50, 20),
                equation="likm,lkjm->lij",
                dtype=tf.float16,
            ),
            param(
                "test 7",
                a_size=(20, 50, 50, 50),
                b_size=(50, 50, 50, 20),
                equation="mlik,lkjm->lij",
                dtype=tf.float16,
            ),
            param(
                "test 8",
                a_size=(2, 5, 5, 5),
                b_size=(5, 5, 5, 2),
                equation="mlik,lkjm",
                dtype=tf.float16,
            ),
            param(
                "test 9",
                a_size=(20, 50, 50, 50),
                b_size=None,
                equation="mlik->imlk",
                dtype=tf.float16,
            ),
            # Activate when cuTENSOR supports it
            # param(
            #     "test 8",
            #     a_size=(20, 50, 50, 50),
            #     b_size=(50, 50, 50, 20),
            #     equation="mlik,lkjm->lij",
            #     dtype=tf.bfloat16,
            # ),
        ]
        # yapf: enable
    )
    def test_einsum_equivalent_results(self,
                                       _,
                                       a_size,
                                       b_size,
                                       equation,
                                       dtype=tf.float32):
        A = tf.random.normal(a_size, dtype=dtype)

        if b_size is not None:
            B = tf.random.normal(b_size, dtype=dtype)

            with tf.GradientTape() as tape:
                tape.watch([A, B])
                tf_native_rslt = tf.einsum(equation, A, B, name="tf_native_einsum")

            tf_native_grads = tape.gradient(tf_native_rslt, [A, B])

            with tf.GradientTape() as tape:
                tape.watch([A, B])
                tf_cutensor_rslt = cutensor.einsum(equation,
                                                   A,
                                                   B,
                                                   name="tf_cuTensor_einsum")

            tf_cutensor_grads = tape.gradient(tf_cutensor_rslt, [A, B])
        else:
            with tf.GradientTape() as tape:
                tape.watch([A])
                tf_native_rslt = tf.einsum(equation, A, name="tf_native_einsum")

            tf_native_grads = tape.gradient(tf_native_rslt, [A])

            with tf.GradientTape() as tape:
                tape.watch([A])
                tf_cutensor_rslt = cutensor.einsum(equation,
                                                   A,
                                                   name="tf_cuTensor_einsum")
            tf_cutensor_grads = tape.gradient(tf_cutensor_rslt, [A])

        self.assertEqual(tf_native_rslt.get_shape(),
                         tf_cutensor_rslt.get_shape())

        self.assertEqual(tf_native_rslt.dtype, tf_cutensor_rslt.dtype)
        self.assertEqual(len(tf_cutensor_grads), len(tf_native_grads))

        self.assertAllClose(tf_native_rslt,
                            tf_cutensor_rslt,
                            rtol=5e-03,
                            atol=5e-03)

        for tf_native_grad, tf_cutensor_grad in zip(tf_native_grads,
                                                    tf_cutensor_grads):
            self.assertAllClose(tf_native_grad,
                                tf_cutensor_grad,
                                rtol=5e-03,
                                atol=5e-03)
            self.assertEqual(tf_native_grad.dtype, tf_cutensor_grad.dtype)

if __name__ == '__main__':
    test.main()


setUp

# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from parameterized import parameterized
from parameterized import param

import tensorflow as tf
from tensorflow.python.platform import test
import tensorflow.test

import cutensor.tensorflow as cutensor

tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

class EinsumcuTENSORTest(tensorflow.test.TestCase):

    def setUp(self):
        super().setUp()
        tf.config.experimental.enable_tensor_float_32_execution(False)

    @parameterized.expand(
        # yapf: disable
        [
            param(
                "test 0",
                a_size=(50, 50),
                b_size=(50, 50),
                equation="ik,kj->ij",
                dtype=tf.float32,
            ),
            param(
                "test 1",
                a_size=(50, 50, 50),
                b_size=(50, 50, 50),
                equation="lik,lkj->lij",
                dtype=tf.float32,
            ),
            param(
                "test 2",


                dtype=tf.float16,
            ),
            param(
                "test 9",
                a_size=(20, 50, 50, 50),
                b_size=None,
                equation="mlik->imlk",
                dtype=tf.float16,
            ),
            # Activate when cuTENSOR supports it
            # param(
            #     "test 8",
            #     a_size=(20, 50, 50, 50),
            #     b_size=(50, 50, 50, 20),
            #     equation="mlik,lkjm->lij",
            #     dtype=tf.bfloat16,
            # ),
        ]
        # yapf: enable
    )
    def test_einsum_equivalent_results(self,
                                       _,
                                       a_size,
                                       b_size,
                                       equation,
                                       dtype=tf.float32):
        A = tf.random.normal(a_size, dtype=dtype)

        if b_size is not None:
            B = tf.random.normal(b_size, dtype=dtype)

            with tf.GradientTape() as tape:
                tape.watch([A, B])
                tf_native_rslt = tf.einsum(equation, A, B, name="tf_native_einsum")

            tf_native_grads = tape.gradient(tf_native_rslt, [A, B])

            with tf.GradientTape() as tape:
                tape.watch([A, B])
                tf_cutensor_rslt = cutensor.einsum(equation,
                                                   A,
                                                   B,
                                                   name="tf_cuTensor_einsum")

            tf_cutensor_grads = tape.gradient(tf_cutensor_rslt, [A, B])
        else:
            with tf.GradientTape() as tape:
                tape.watch([A])
                tf_native_rslt = tf.einsum(equation, A, name="tf_native_einsum")

            tf_native_grads = tape.gradient(tf_native_rslt, [A])

            with tf.GradientTape() as tape:
                tape.watch([A])
                tf_cutensor_rslt = cutensor.einsum(equation,
                                                   A,
                                                   name="tf_cuTensor_einsum")
            tf_cutensor_grads = tape.gradient(tf_cutensor_rslt, [A])

        self.assertEqual(tf_native_rslt.get_shape(),
                         tf_cutensor_rslt.get_shape())

        self.assertEqual(tf_native_rslt.dtype, tf_cutensor_rslt.dtype)
        self.assertEqual(len(tf_cutensor_grads), len(tf_native_grads))

        self.assertAllClose(tf_native_rslt,
                            tf_cutensor_rslt,
                            rtol=5e-03,
                            atol=5e-03)

        for tf_native_grad, tf_cutensor_grad in zip(tf_native_grads,
                                                    tf_cutensor_grads):
            self.assertAllClose(tf_native_grad,
                                tf_cutensor_grad,
                                rtol=5e-03,
                                atol=5e-03)
            self.assertEqual(tf_native_grad.dtype, tf_cutensor_grad.dtype)

if __name__ == '__main__':
    test.main()


# ! /usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name(s) of the copyright holder(s) nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import tensorflow as tf
from tensorflow.python.framework import ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import special_math_ops
from tensorflow.python.framework.load_library import load_op_library

from ..common import normalize_subscript

import glob
import os
pattern = os.path.join(os.path.dirname(__file__), 'binding*.so')
glob_res = glob.glob(pattern)
binding_file, = glob_res
einsum_lib = tf.load_op_library(binding_file)

def einsum(equation, *inputs, **kwargs):
    name = kwargs.pop('name', None)

    if kwargs:
        raise TypeError(
            'invalid keyword arguments for this function: ' +
            ', '.join([format(key) for key in sorted(list(kwargs.keys()))]))

    with ops.name_scope(name, 'einsum', [equation, inputs]):
        inputs = list(inputs)

        input_shapes = [x.get_shape() for x in inputs]
        input_axis_labels, output_axis_labels = special_math_ops._einsum_v1_parse_and_resolve_equation(
            equation, input_shapes)

        axis_labels = set(''.join(input_axis_labels) + output_axis_labels)

        for a in axis_labels:
            for input_labels in input_axis_labels:
                if (len(input_axis_labels) == 1 and
                        input_labels.count(a) == 2 and
                        input_labels == input_labels[::-1] and
                        '->' not in equation):
                    return math_ops.trace(inputs[0])
                if input_labels.count(a) > 1:
                    raise ValueError(
                        'Subscript not supported: an axis appears more than once: %s'
                        % input_labels)

        for a in axis_labels:

            input_count = sum(1 for s in input_axis_labels if a in s)

            if input_count > 2 and a not in output_axis_labels:
                tf.compat.v1.logging.warn(
                    'Falling back to exponential-space implementation of einsum()'
                    ' because index "%s" is summed over more than two inputs.',
                    a)
                return special_math_ops._exponential_space_einsum(
                    equation, *inputs)

        equation = ','.join(input_axis_labels) + '->' + output_axis_labels
        if len(inputs) == 1:
            # inputs.append(inputs[0])
            inputs.append(tf.constant([0], dtype=inputs[0].dtype))
        return einsum_lib.einsum_cu_tensor(input_0=inputs[0],
                                           input_1=inputs[1],
                                           equation=equation)

@ops.RegisterGradient("EinsumCuTensor")
def _einsum_cu_tensor_grad(op, grad):
    A = op.inputs[0]
    B = op.inputs[1]

    subscript, _ = normalize_subscript(op.get_attr("equation").decode())
    lhs, modeC = subscript.split('->')
    if ',' in lhs:
        modeA, modeB = lhs.split(',')

        grad_A = einsum_lib.einsum_cu_tensor(input_0=grad,
                                             input_1=B,
                                             equation=modeC + ',' + modeB + '->' +
                                             modeA)

        grad_B = einsum_lib.einsum_cu_tensor(input_0=A,
                                             input_1=grad,
                                             equation=modeA + ',' + modeC + '->' +
                                             modeB)

        return [grad_A, grad_B]
    else:
        grad = einsum_lib.einsum_cu_tensor(input_0=grad,
                                           input_1=B,
                                           equation=modeC + '->' + lhs)
        return [grad, B]


einsum

# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import tensorflow as tf
from tensorflow.python.framework import ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import special_math_ops
from tensorflow.python.framework.load_library import load_op_library

from ..common import normalize_subscript

import glob
import os
pattern = os.path.join(os.path.dirname(__file__), 'binding*.so')
glob_res = glob.glob(pattern)
binding_file, = glob_res
einsum_lib = tf.load_op_library(binding_file)

def einsum(equation, *inputs, **kwargs):
    name = kwargs.pop('name', None)

    if kwargs:
        raise TypeError(
            'invalid keyword arguments for this function: ' +
            ', '.join([format(key) for key in sorted(list(kwargs.keys()))]))

    with ops.name_scope(name, 'einsum', [equation, inputs]):
        inputs = list(inputs)

        input_shapes = [x.get_shape() for x in inputs]
        input_axis_labels, output_axis_labels = special_math_ops._einsum_v1_parse_and_resolve_equation(
            equation, input_shapes)

        axis_labels = set(''.join(input_axis_labels) + output_axis_labels)

        for a in axis_labels:
            for input_labels in input_axis_labels:
                if (len(input_axis_labels) == 1 and
                        input_labels.count(a) == 2 and
                        input_labels == input_labels[::-1] and
                        '->' not in equation):
                    return math_ops.trace(inputs[0])
                if input_labels.count(a) > 1:
                    raise ValueError(
                        'Subscript not supported: an axis appears more than once: %s'
                        % input_labels)

        for a in axis_labels:

            input_count = sum(1 for s in input_axis_labels if a in s)

            if input_count > 2 and a not in output_axis_labels:
                tf.compat.v1.logging.warn(
                    'Falling back to exponential-space implementation of einsum()'
                    ' because index "%s" is summed over more than two inputs.',
                    a)
                return special_math_ops._exponential_space_einsum(
                    equation, *inputs)

        equation = ','.join(input_axis_labels) + '->' + output_axis_labels
        if len(inputs) == 1:
            # inputs.append(inputs[0])
            inputs.append(tf.constant([0], dtype=inputs[0].dtype))
        return einsum_lib.einsum_cu_tensor(input_0=inputs[0],
                                           input_1=inputs[1],
                                           equation=equation)

@ops.RegisterGradient("EinsumCuTensor")
def _einsum_cu_tensor_grad(op, grad):
    A = op.inputs[0]
    B = op.inputs[1]

    subscript, _ = normalize_subscript(op.get_attr("equation").decode())
    lhs, modeC = subscript.split('->')
    if ',' in lhs:
        modeA, modeB = lhs.split(',')

        grad_A = einsum_lib.einsum_cu_tensor(input_0=grad,
                                             input_1=B,
                                             equation=modeC + ',' + modeB + '->' +
                                             modeA)

        grad_B = einsum_lib.einsum_cu_tensor(input_0=A,
                                             input_1=grad,
                                             equation=modeA + ',' + modeC + '->' +


            input_count = sum(1 for s in input_axis_labels if a in s)

            if input_count > 2 and a not in output_axis_labels:
                tf.compat.v1.logging.warn(
                    'Falling back to exponential-space implementation of einsum()'
                    ' because index "%s" is summed over more than two inputs.',
                    a)
                return special_math_ops._exponential_space_einsum(
                    equation, *inputs)

        equation = ','.join(input_axis_labels) + '->' + output_axis_labels
        if len(inputs) == 1:
            # inputs.append(inputs[0])
            inputs.append(tf.constant([0], dtype=inputs[0].dtype))
        return einsum_lib.einsum_cu_tensor(input_0=inputs[0],
                                           input_1=inputs[1],
                                           equation=equation)

@ops.RegisterGradient("EinsumCuTensor")
def _einsum_cu_tensor_grad(op, grad):
    A = op.inputs[0]
    B = op.inputs[1]

    subscript, _ = normalize_subscript(op.get_attr("equation").decode())
    lhs, modeC = subscript.split('->')
    if ',' in lhs:
        modeA, modeB = lhs.split(',')

        grad_A = einsum_lib.einsum_cu_tensor(input_0=grad,
                                             input_1=B,
                                             equation=modeC + ',' + modeB + '->' +
                                             modeA)

        grad_B = einsum_lib.einsum_cu_tensor(input_0=A,
                                             input_1=grad,
                                             equation=modeA + ',' + modeC + '->' +
                                             modeB)

        return [grad_A, grad_B]
    else:
        grad = einsum_lib.einsum_cu_tensor(input_0=grad,
                                           input_1=B,
                                           equation=modeC + '->' + lhs)
        return [grad, B]


#!/usr/bin/env python3

# SPDX-FileCopyrightText: Copyright (c) 2018-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import struct
import argparse
import numpy
import warnings

delimiter = ','

def fixInput(val):
    if val is None:
        return 0
    try:
        retval = int(val)
    except ValueError:
        retval = ord(val)

    return retval

if len(sys.argv) != 5 and len(sys.argv) != 6:
    print("Usage:")
    print("\tcsv_to_binary.py <input filename> <column choice> <datatype> <output filename> [delimiter]")
    print()
    print("This program converts one column of a text file containing a table of data,")
    print("(a comma-separated values file by default), into a binary file.")
    print()
    print("The <column choice> should be an integer in the range [0, N-1], where N is the number of columns.")
    print("The <datatype> option should be one of 'int', 'long', 'float', 'double', or 'string'.")
    print("'string' keeps the text, converting it to UTF-16 with no separators between the values.")
    print("The [delimiter] is an optional argument, and defaults to '%s'" % delimiter)
    print("Some delimiters may need to be surrounded by quotation marks or prefixed by a backslash, depending on")
    print("the shell, for example space, semicolon, or vertical pipe, due to the command line parsing")
    print("interpreting the space or semicolon as a parameter separator or command separator, instead of a")
    print("parameter to this script.")
    print()
    print("Examples:")
    print("    text_to_binary.py ExampleFloatData.csv 2 float ZValues.bin")
    print("    text_to_binary.py ExampleTable.txt 5 long Dates.bin '|'")
    print("    text_to_binary.py SpaceSeparatedData.txt 0 int FirstColumn.bin ' '")
    print()
    exit()

in_fname = sys.argv[1]
col_num = sys.argv[2]
datatype = sys.argv[3]
out_fname = sys.argv[4]

if len(sys.argv) == 6:
    delimiter = sys.argv[5]

# Add more datatypes if needed
if datatype == "int":
    dtype = "int32"
elif datatype == "long":
    dtype = "int64"
elif datatype == "float":
    dtype = "float32"
elif datatype == "double":
    dtype = "float64"
elif datatype == "string":
    dtype = "str"
else:
    print("Please select datatype int, long, float, double, or string")
    exit()

print("Reading column " + col_num + ", of type " + datatype + "...")

chunk_size = 10000000
finished = False
offset = 0
with open(str(in_fname), "r") as inFile:
    with open(str(out_fname), "wb") as newFile:
        with warnings.catch_warnings():
            while not finished:
                in_data=numpy.genfromtxt(inFile, dtype=dtype,
                max_rows=chunk_size, usecols=(int(col_num),), delimiter=delimiter, loose=False)

                if offset == 0:
                    # don't warn about an empty file after we have read something
                    warnings.filterwarnings('ignore', r'genfromtxt: Empty input file:')

                if in_data.size > 0:
                    in_data.tofile(newFile)
                    offset += in_data.size
                else:
                    finished = True
if offset != 0:
    print('Wrote '+str(offset)+' '+datatype+'s to '+str(out_fname))
else:
    print('Wrote no data')


fixInput

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import struct
import argparse
import numpy
import warnings

delimiter = ','

def fixInput(val):
    if val is None:
        return 0
    try:
        retval = int(val)
    except ValueError:
        retval = ord(val)

    return retval

if len(sys.argv) != 5 and len(sys.argv) != 6:
    print("Usage:")
    print("\tcsv_to_binary.py <input filename> <column choice> <datatype> <output filename> [delimiter]")
    print()
    print("This program converts one column of a text file containing a table of data,")
    print("(a comma-separated values file by default), into a binary file.")
    print()
    print("The <column choice> should be an integer in the range [0, N-1], where N is the number of columns.")
    print("The <datatype> option should be one of 'int', 'long', 'float', 'double', or 'string'.")
    print("'string' keeps the text, converting it to UTF-16 with no separators between the values.")
    print("The [delimiter] is an optional argument, and defaults to '%s'" % delimiter)
    print("Some delimiters may need to be surrounded by quotation marks or prefixed by a backslash, depending on")
    print("the shell, for example space, semicolon, or vertical pipe, due to the command line parsing")
    print("interpreting the space or semicolon as a parameter separator or command separator, instead of a")
    print("parameter to this script.")
    print()
    print("Examples:")
    print("    text_to_binary.py ExampleFloatData.csv 2 float ZValues.bin")