Define accelerated blocks

Warning

I’m not satisfied by the syntax for Pythran blocks so I (PA) proposed an alternative syntax in issue #6.

Transonic blocks can be used with classes and more generally in functions with lines that cannot be compiled by Pythran.

import numpy as np

from transonic import Transonic

ts = Transonic()

# don't define classes in Pythran file! Here, no problem...


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

    def compute(self, n):

        a = self.a
        b = self.b

        if ts.is_transpiled:
            result = ts.use_block("block0")
        else:
            # transonic block (
            #     float[][] a, b;
            #     int n
            # )
            # transonic block (
            #     float[][][] a, b;
            #     int n
            # )
            result = np.zeros_like(a)
            for _ in range(n):
                result += a**2 + b**3

        return result


if __name__ == "__main__":

    shape = 100, 100
    a = np.random.rand(*shape)
    b = np.random.rand(*shape)

    obj = MyClass(a, b)

    obj.compute(10)

    if ts.is_transpiled:
        ret = obj.compute(10)
        ts.is_transpiled = False
        ret1 = obj.compute(10)
        ts.is_transpiled = True
        assert np.allclose(ret, ret1)
        print("allclose OK")

For blocks, we need a little bit more of Python.

• At import time, we have ts = Transonic(), which detects which Pythran module should be used and imports it. This is done at import time since we want to be very fast at run time.

• In the function, we define a block with three lines of Python and special Pythran annotations (# transonic block). The 3 lines of Python are used (i) at run time to choose between the two branches (is_transpiled or not) and (ii) at compile time to detect the blocks.

Note that the annotations in the command # transonic block are different (and somehow easier to write) than in the standard command # pythran export.

Warning

The two branches of the if ts.is_transpiled are not equivalent! The user has to be careful because it is not difficult to write such buggy code:

c = 0
if ts.is_transpiled:
    a, b = ts.use_block("buggy_block")
else:
    # transonic block ()
    a = b = c = 1

assert c == 1

Note

The Pythran keyword or cannot be used in block annotations (not yet implemented, see issue #11).

Blocks can now also be defined with type hints!

import numpy as np

from transonic import Array, NDim, Transonic, Type

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

        if ts.is_transpiled:
            result = ts.use_block("block0")
        else:
            # transonic block (A a, b; A1 c; int n)
            # transonic block (int a, b, c; float n)

            result = np.zeros_like(a)
            for _ in range(n):
                result += a**2 + b**3

        return result


if __name__ == "__main__":

    shape = 100, 100
    a = np.random.rand(*shape)
    b = np.random.rand(*shape)

    obj = MyClass(a, b)

    obj.compute(10)

    if ts.is_transpiled:
        ret = obj.compute(10)
        ts.is_transpiled = False
        ret1 = obj.compute(10)
        ts.is_transpiled = True
        assert np.allclose(ret, ret1)
        print("allclose OK")