⛳

get_ipython().run_line_magic()について調べてみた。

2024/06/30に公開

JAX

IPython

マジックコマンド

tech

はじめに

Pythonでjaxについて勉強している際に

    get_ipython().run_line_magic('timeit', f'-n {repeat} numpy_mod(size)')

このコードを目の当たりしましたが、この意味が全然分からなかったので、調べてみました。
今回使うコードは以下のようなものです。

import jax.numpy as jnp
from jax import jit
from functools import partial

# (size, size)の行列を作ってMod計算
@partial(jit, static_argnums=(0,))
def jax_jit_mod(size):
    x = jnp.arange(size, dtype=jnp.int32)
    mat = x[None, :] * x[:, None] # (size, size)
    return mat % 256

for i in range(4):
    size = 10**(i+1)
    repeat = 10**(4-i)
    print("size =", size, "repeat =", repeat)
    get_ipython().run_line_magic('timeit', f'-n {repeat} jax_jit_mod(size).block_until_ready()') # jitありJAX

コードの解説

全体の流れ

全体の流れとしては、異なるサイズの行列に対する計算速度を測定する際にjaxのjitを用いて高速化してみたものです。行列計算はxの外積を計算し、各要素を256で割るという処理をしています。

jitの説明

jitを使って関数をJITコンパイルします。

JIT（Just-In-Time）コンパイルは、プログラムの一部を実行時にネイティブマシンコードに変換します。これにより、通常のインタープリタ実行よりも高速に動作します。
特に、数値計算や行列操作などの重い計算に対して、JITコンパイルによる最適化は大きなパフォーマンス向上をもたらします。

partial は、関数の一部の引数やキーワード引数を事前に設定して新しい関数を作成するためのものです。JAXのjitデコレータと一緒に使うことで、特定の引数を静的に指定できます。

@partial(jit, static_argnums=(0,)) の意味

@partial(jit, static_argnums=(0,)) は、以下のように解釈できます：

jit デコレータで関数をJITコンパイルする。
static_argnums=(0,) で、関数の最初の引数（0番目の引数）を静的引数として指定する。

get_ipython().run_line_magic('timeit', f'-n {repeat} jax_jit_mod(size).block_until_ready()') の説明

get_ipython().run_line_magic('timeit', f'-n {repeat} numpy_mod(size)') は、IPython環境でタイミング計測を行うためのコードです。具体的には、%timeit マジックコマンドを使って、関数の実行時間を測定します。以下に詳しく説明します。

IPython マジックコマンド

IPythonには、特定のタスクを簡略化するためのマジックコマンドが多数用意されています。その中でも %timeit は、コードの実行時間を測定するために使われます。

timeitマジックコマンド

%timeit は、指定したコードの実行時間を測定し、複数回の実行結果の平均を表示します。これにより、コードの性能を評価するのに役立ちます。