Nearly as with Pythran

This example shows how to use Transonic nearly how you would use Pythran.

import numpy as np

# don't import any random modules in a Pythran file. Here, no problem!
from fluiddyn.util import mpi

from transonic import boost

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


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


if __name__ == "__main__":

    n0, n1 = 100, 200
    a0 = np.random.rand(n0, n1)
    a1 = np.random.rand(n0, n1)

    result = func(a0, a1)
    if mpi.nb_proc > 1:
        result = mpi.comm.allreduce(result, op=mpi.MPI.MAX)
    mpi.printby0(result)

    a0 = (1000 * a0).astype(int)
    result = func(a0, a1)
    if mpi.nb_proc > 1:
        result = mpi.comm.allreduce(result, op=mpi.MPI.MAX)
    mpi.printby0(result)

Most of this code looks familiar to Pythran users. The differences:

One can use (for example) h5py and mpi4py (of course not in the Pythran functions).
# transonic def instead of # pythran export.
A tiny bit of Python… The decorator @boost replaces the Python function by the pythranized function if Transonic has been used to produced the associated Pythran file.

Function calls

Boosted functions support function calls, including imported functions from a Python file in the same directory!

import numpy as np
from local_module import multiply

from transonic import boost

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


def my_log(b):
    return np.log(b)


@boost
def func(a, b):
    c = multiply(a, b)
    return (c * my_log(b)).max()

Note that no Pythran signatures or annotations are needed for functions called in boosted functions.