<jupyter_start><jupyter_text>IPython + jit decoratorLet's say that we have this kind of code:<jupyter_code>import numpy as np
image = np.random.rand(1024, 1024)
def laplace(image):
    """Laplace operator in NumPy for 2D images."""
    laplacian = (
        image[:-2, 1:-1] + image[2:, 1:-1] + image[1:-1, :-2] + image[1:-1, 2:]
        - 4*image[1:-1, 1:-1]
    )
    thresh = np.abs(laplacian) > 0.05
    return thresh<jupyter_output><empty_output><jupyter_text>We can define a jitted function in IPython (and Jupyter, which uses IPython). The only limitation is that everything needed for the function (even `np`) has to be (re)defined in the same cell!<jupyter_code>import numpy as np
from transonic import jit

@jit
def laplace_pythran(image):
    """Laplace operator in NumPy for 2D images."""
    laplacian = (
        image[:-2, 1:-1] + image[2:, 1:-1] + image[1:-1, :-2] + image[1:-1, 2:]
        - 4*image[1:-1, 1:-1]
    )
    thresh = np.abs(laplacian) > 0.05
    return thresh<jupyter_output><empty_output><jupyter_text>We first call the function to launch the compilation ("warmup the jit"):<jupyter_code>result = laplace_pythran(image)<jupyter_output><empty_output><jupyter_text>Since the compilation is quit long, we explicitelly wait for the extension to benchmark the compiled version.<jupyter_code>from transonic import wait_for_all_extensions
wait_for_all_extensions()<jupyter_output><empty_output><jupyter_text>Now, let's benchmark the 2 functions:<jupyter_code>from transonic.util import timeit_verbose

namespace = {"laplace": laplace, "laplace_pythran": laplace_pythran, "image": image}

norm = timeit_verbose("laplace(image)", globals=namespace)
time_pythran = timeit_verbose("laplace_pythran(image)", globals=namespace, norm=norm)
print(f"\nIt corresponds to a speedup of ~ {norm/time_pythran:.1f}!")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12870220786214323, 'fraction_numerical': 0.028002154011847066, 'mean_word_length': 5.193333333333333, 'pattern_counts': {'":': 3, '<': 21, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '574042', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 700, 'n_tokens_neox': 654, 'n_words': 168, 'path': 'compile_jit.ipynb', 'repo_name': 'fluiddyn/transonic-demos', 'size': '1857'}