Accelerated functions with type hints

First a minimalist example:

# don't import h5py and mpi4py in a Pythran file, here, no problem!
import h5py
import mpi4py

from transonic import boost


@boost
def myfunc(a: int, b: float):
    return a * b

Nice but very limited… So it possible to mix type hints and # pythran def commands.

import numpy as np

# don't import skimage in a Pythran file. Here, no problem!
from skimage.filters import sobel

from transonic import boost

# transonic def func(float[][], float[][])
# transonic def func(int[][], float[][])


@boost
def func(a: float, b: float):
    return (a * np.log(b)).max()


@boost
def func1(a: int, b: float):
    return a * np.cos(b)

Moreover, you can write:

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

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

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


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

Yes, this one is neat!

Note

Note that the array types A and A1 are bound. When A.ndim == 1, then A1.ndim == 2, and when A.ndim == 3, then A1.ndim == 4. No signatures are produced for A.ndim == 1, A1.ndim == 4.

Note that one can also just write Pythran type-string in type annotations:

@boost
def myfunc(a: "float[3, :]", b: float):
    ...

Array types with only one number of dimension can simply be defined like this:

from transonic import Array

A = Array[float, "3d"]

Which has actually the same effect as:

A = "float[:, :, :]"

But, one can also specify the memory layout, for example for a C-order array and a strided array:

A = Array[float, "3d", "C"]
A_strided = Array[float, "3d", "strided"]

Oh, and you can also write:

import numpy as np
from transonic import str2type, typeof

T = str2type("(int, float)")  # a tuple
A = typeof(np.empty((2, 2)))

There is also an Union “type” that can be used similarly as typing.Union:

from transonic import Array, Union

U = Union[float, Array[float, "3d"]]

All Transonic types are defined in the module transonic.typing.