How to use the transonic.NDim function in transonic

from functools import partial

import numpy as np

import transonic as ts
from transonic import Type, NDim, Array, Union

T = Type(int, np.complex128)
N = NDim(1, 3)

A = Array[T, N]
A1 = Array[np.float32, N + 1]

A3d = Array[np.float32, "3d"]
N1 = NDim(4, 5)
N1 = NDim(4, 5)

T = Type(int, np.complex128)


@ts.boost
def compute(a: A, b: A, c: T, d: Union[A, A1], e: str):
    print(e)
    tmp = a + b
    if 1 and 2:
        tmp *= 2
    return tmp


main = partial(lambda x: x, lambda x: x)

import numpy as np
import transonic as ts
from transonic import Type, NDim, Array, Union

T = Type("T")
T1 = Type("T1")
N = NDim("N")

A = Array[T, N]
A1 = Array[T1, N + 1]

# for coverage
assert repr(N - 1) == "N - 1"
print(repr(A1))


repr(Union[A, A1])


@ts.boost
def compute(a: A, b: A1, c: T, d: A, e: str):
    print(e)
    tmp = a + b

import numpy as np

from transonic import Transonic, Type, NDim, Array

T = Type(float, complex)
N = NDim(1, 2)
A = Array[T, N]
A1 = Array[T, N + 1]

ts = Transonic()


class MyClass:
    def __init__(self, a, b):
        self.a = a
        self.b = b

    def compute(self, n):

        a = self.a
        b = self.b

from transonic.log import logger
from transonic.util import get_source_without_decorator, format_str
from transonic.annotation import compute_signatures_from_typeobjects


def trans_def_method(func):
    func.__transonic__ = "trans_def_method"
    return func


# logger.set_level("debug")


from transonic import Array, Type, NDim

A = Array[Type(float, int), NDim(1, 2)]


class Transmitter:
    freq: float

    def __init__(self):
        pass

    @trans_def_method
    def __call__(self, inp: A):
        """My docstring"""
        return inp * np.exp(np.arange(len(inp)) * self.freq * 1j)

    @trans_def_method
    def call_with_print(self, inp: A):
        """call + print"""

import transonic as ts
from transonic import Type, NDim, Array, Union

import numpy as np
import skimage

T = Type(int, np.complex128)

dim = 2
dim += 1

N = NDim(1, dim)

A = Array[T, N]
A1 = Array[np.float32, N + 1]

A3d = Array[np.float32, "3d"]
N1 = NDim(4, 5)
N1 = NDim(4, 5)

T = Type(int, np.complex128)

a_type_var = "hello"
myconst = 0

cdict = skimage.color.color_dict

@ts.boost

import numpy as np
import transonic as ts
from transonic import Type, NDim, Array

T = Type("T")
N = NDim("N")

A = Array[T, N]
A1 = Array[np.int8, N + 1]


@ts.boost
def compute(a: A, b: A, c: T, d: A1, e: str):
    print(e)
    tmp = a + b
    return tmp


for dtype in [int, np.complex128]:
    for ndim in [1, 3]:
        ts.make_signature(compute, T=dtype, N=ndim)

import numpy as np

from transonic import Type, NDim, Array, boost

T = Type(np.float64, np.complex128)
N = NDim(1)
A = Array[T, N]


@boost
def func(a: A):
    i: int
    n: int = a.shape[0]

    for i in range(n):
        a[i] = a[i] + 1.

from transonic import boost, Type, NDim, Shape, Array

T = Type(int, float)

# here the shape of the array is only defined with the ShapeVar
A = Array[T, Shape("[3, :]", "[3, :, :]", "[::, ::]", "[::, ::, ::]")]


@boost
def compute(a: A, b: A, c: T):
    return a + b


# if there is a NDimVar, we can use the ellipsis
A1 = Array[T, NDim(1, 3), Shape("[3, ...]", "[::, ...]")]


@boost
def compute1(a: A1, b: A1, c: T):
    return c * (a + b)

import numpy as np

from transonic import Transonic, Type, NDim, Array

T = Type(float, complex)
N = NDim(2, 3)
A = Array[T, N]
A1 = Array[T, N + 1]

ts = Transonic()


class MyClass:
    def __init__(self, a, b):
        self.a = a
        self.b = b

    def compute(self, n):

        a = self.a
        b = self.b

import numpy as np
import skimage

T = Type(int, np.complex128)

dim = 2
dim += 1

N = NDim(1, dim)

A = Array[T, N]
A1 = Array[np.float32, N + 1]

A3d = Array[np.float32, "3d"]
N1 = NDim(4, 5)
N1 = NDim(4, 5)

T = Type(int, np.complex128)

a_type_var = "hello"
myconst = 0

cdict = skimage.color.color_dict

@ts.boost
def compute(a: A, b: A, c: T, d: Union[A, A1], e: str):
    print(e)
    tmp = a + b + myconst
    return tmp