LivebookもEXLAを使ってスピードアップしたい！

Notebook dependences and setupの欄に次のように記述します

Mix.install([
  {:nx, "~> 0.3.0"},
  {:exla, "~> 0.3.0"},
],
config: [
          nx: [
            default_backend: EXLA.Backend,
            default_defn_options: [compiler: EXLA],
          ]
        ],
system_env: [
    XLA_TARGET: "cuda111"
  ],
force: false # XLA_TARGETの変更で、exlaを再コンパイルしたい時はtrueにする
)

XLAターゲットの指定

XLA_TARGETの値	ターゲット
cpu	cpu (default)
cuda111	cuda version 11.x
rocm	AMD GPU
tpu	libtpu

2022/10/29追記
System.put_env("XLA_TARGET", "cuda111")をNotebook dependences and setupの欄の最初に記述する事で指定できます。

コメントいただいた方法で、うまくいきました。ありがとうございました。

2022/1/2
Mix.installの第二引数のoptsにsystem_env:　で環境変数を指定できました。これでXLA_TARGETの指定できるとわかりました。この方法に変更しました。

XLA_TARGET変更時の注意

XLA_TARGETを指定しないで実行すると、XLA_TARGET=cpuでexlaのコンパイルが実行されます。
この結果は、
~/.cache/mix/installs/elixir-1.14.0-erts-13.0.4/4bd451fe49c8c67a1da2c1482de3aaa1/deps/exla/cache
にキャッシュされます。

setupの、XLA_TARGETをcuda111に変更しても、キャッシュの影響で、exlaのcompileは再実行されない為、もCPUのままで、GPUが使われません。

`~/.cache/mix/installs/elixir-1.14.0-erts-13.0.4/4bd451fe49c8c67a1da2c1482de3aaa1/
を削除することで、XLA_TARAGETの値を反映した状態で、exlaのcompileが再実行されました。

※elixir-1.14.0-erts-13.0.4/4bd451fe49c8c67a1da2c1482de3aaa1の部分は、実行環境によって異なります。

Mix.installの第二引数以降に、force: trueを付ける事で、exlaのcompileが再実行され、XLA_TARGETの値が反映されます。

参考
https://hexdocs.pm/exla/EXLA.html#module-configuration
https://hexdocs.pm/mix/Mix.html#install/2