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
.