Skip to content

Instantly share code, notes, and snippets.

@sabiwara
Last active May 29, 2022 22:18
Show Gist options
  • Save sabiwara/589eeccbec0932c85faed970119be4dc to your computer and use it in GitHub Desktop.
Save sabiwara/589eeccbec0932c85faed970119be4dc to your computer and use it in GitHub Desktop.
Nx hands On Livebook - TokyoEx 2022/05/29

Nx hands On Livebook - TokyoEx 2022/05/29

Mix.install([
  {:nx, "~> 0.2.0"},
  {:exla, "~> 0.2.0"}
])
Resolving Hex dependencies...
Dependency resolution completed:
New:
  complex 0.4.1
  elixir_make 0.6.3
  exla 0.2.0
  nx 0.2.0
  xla 0.3.0
* Getting nx (Hex package)
* Getting exla (Hex package)
* Getting elixir_make (Hex package)
* Getting xla (Hex package)
* Getting complex (Hex package)
==> complex
Compiling 2 files (.ex)
Generated complex app
==> nx
Compiling 24 files (.ex)
Generated nx app
==> elixir_make
Compiling 1 file (.ex)
Generated elixir_make app
==> xla
Compiling 2 files (.ex)
Generated xla app

14:17:11.910 [info]  Found a matching archive (xla_extension-aarch64-darwin-cpu.tar.gz), going to download it

14:17:20.861 [info]  Successfully downloaded the XLA archive
==> exla
Unpacking /Users/sabiwara/Library/Caches/xla/0.3.0/cache/download/xla_extension-aarch64-darwin-cpu.tar.gz into /Users/sabiwara/Library/Caches/mix/installs/elixir-1.13.4-erts-13.0/58b1d08ea5c0d2502d3eea455c1891c3/deps/exla/cache
c++ -fPIC -I/Users/sabiwara/.asdf/installs/erlang/25.0/erts-13.0/include -Icache/xla_extension/include -O3 -Wall -Wno-sign-compare -Wno-unused-parameter -Wno-missing-field-initializers -Wno-comment -shared -std=c++14 c_src/exla/exla.cc c_src/exla/exla_nif_util.cc c_src/exla/exla_client.cc -o cache/libexla.so -Lcache/xla_extension/lib -lxla_extension -flat_namespace -undefined suppress
In file included from c_src/exla/exla.cc:4:
In file included from c_src/exla/exla_client.h:12:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/gpu_device.h:22:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h:34:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/client/local_client.h:27:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/compiler.h:30:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_assignment.h:28:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/heap_simulator.h:29:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/hlo_alias_analysis.h:25:
cache/xla_extension/include/tensorflow/compiler/xla/service/hlo_buffer.h:100:46: warning: 'color' is deprecated: Use Layout::memory_space instead. [-Wdeprecated-declarations]
    BufferValue::Color result = values()[0]->color();
                                             ^
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_value.h:111:3: note: 'color' has been explicitly marked deprecated here
  ABSL_DEPRECATED("Use Layout::memory_space instead.")
  ^
cache/xla_extension/include/absl/base/attributes.h:648:49: note: expanded from macro 'ABSL_DEPRECATED'
#define ABSL_DEPRECATED(message) __attribute__((deprecated(message)))
                                                ^
In file included from c_src/exla/exla.cc:4:
In file included from c_src/exla/exla_client.h:12:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/gpu_device.h:22:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h:34:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/client/local_client.h:27:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/compiler.h:30:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_assignment.h:28:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/heap_simulator.h:29:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/hlo_alias_analysis.h:25:
cache/xla_extension/include/tensorflow/compiler/xla/service/hlo_buffer.h:102:32: warning: 'color' is deprecated: Use Layout::memory_space instead. [-Wdeprecated-declarations]
      DCHECK_EQ(result, value->color());
                               ^
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_value.h:111:3: note: 'color' has been explicitly marked deprecated here
  ABSL_DEPRECATED("Use Layout::memory_space instead.")
  ^
cache/xla_extension/include/absl/base/attributes.h:648:49: note: expanded from macro 'ABSL_DEPRECATED'
#define ABSL_DEPRECATED(message) __attribute__((deprecated(message)))
                                                ^
In file included from c_src/exla/exla.cc:4:
In file included from c_src/exla/exla_client.h:12:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/gpu_device.h:22:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h:34:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/client/local_client.h:27:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/compiler.h:30:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_assignment.h:37:
cache/xla_extension/include/tensorflow/compiler/xla/service/memory_space_assignment.h:1000:1: warning: 'Options' defined as a struct here but previously declared as a class; this is valid, but may result in linker errors under the Microsoft C++ ABI [-Wmismatched-tags]
struct Options {
^
cache/xla_extension/include/tensorflow/compiler/xla/service/memory_space_assignment.h:37:1: note: did you mean struct here?
class Options;
^~~~~
struct
In file included from c_src/exla/exla.cc:4:
In file included from c_src/exla/exla_client.h:12:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/gpu_device.h:22:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h:34:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/client/local_client.h:27:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/compiler.h:30:
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_assignment.h:600:18: warning: 'set_color' is deprecated: Use Layout::memory_space instead. [-Wdeprecated-declarations]
          value->set_color(BufferValue::Color(
                 ^
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_value.h:118:3: note: 'set_color' has been explicitly marked deprecated here
  ABSL_DEPRECATED("Use Layout::memory_space instead.")
  ^
cache/xla_extension/include/absl/base/attributes.h:648:49: note: expanded from macro 'ABSL_DEPRECATED'
#define ABSL_DEPRECATED(message) __attribute__((deprecated(message)))
                                                ^
In file included from c_src/exla/exla.cc:4:
In file included from c_src/exla/exla_client.h:12:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/gpu_device.h:22:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h:34:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/client/local_client.h:27:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/compiler.h:30:
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_assignment.h:603:18: warning: 'set_color' is deprecated: Use Layout::memory_space instead. [-Wdeprecated-declarations]
          value->set_color(BufferValue::Color(0));
                 ^
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_value.h:118:3: note: 'set_color' has been explicitly marked deprecated here
  ABSL_DEPRECATED("Use Layout::memory_space instead.")
  ^
cache/xla_extension/include/absl/base/attributes.h:648:49: note: expanded from macro 'ABSL_DEPRECATED'
#define ABSL_DEPRECATED(message) __attribute__((deprecated(message)))
                                                ^
5 warnings generated.
In file included from c_src/exla/exla_client.cc:1:
In file included from c_src/exla/exla_client.h:12:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/gpu_device.h:22:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h:34:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/client/local_client.h:27:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/compiler.h:30:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_assignment.h:28:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/heap_simulator.h:29:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/hlo_alias_analysis.h:25:
cache/xla_extension/include/tensorflow/compiler/xla/service/hlo_buffer.h:100:46: warning: 'color' is deprecated: Use Layout::memory_space instead. [-Wdeprecated-declarations]
    BufferValue::Color result = values()[0]->color();
                                             ^
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_value.h:111:3: note: 'color' has been explicitly marked deprecated here
  ABSL_DEPRECATED("Use Layout::memory_space instead.")
  ^
cache/xla_extension/include/absl/base/attributes.h:648:49: note: expanded from macro 'ABSL_DEPRECATED'
#define ABSL_DEPRECATED(message) __attribute__((deprecated(message)))
                                                ^
In file included from c_src/exla/exla_client.cc:1:
In file included from c_src/exla/exla_client.h:12:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/gpu_device.h:22:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h:34:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/client/local_client.h:27:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/compiler.h:30:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_assignment.h:28:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/heap_simulator.h:29:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/hlo_alias_analysis.h:25:
cache/xla_extension/include/tensorflow/compiler/xla/service/hlo_buffer.h:102:32: warning: 'color' is deprecated: Use Layout::memory_space instead. [-Wdeprecated-declarations]
      DCHECK_EQ(result, value->color());
                               ^
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_value.h:111:3: note: 'color' has been explicitly marked deprecated here
  ABSL_DEPRECATED("Use Layout::memory_space instead.")
  ^
cache/xla_extension/include/absl/base/attributes.h:648:49: note: expanded from macro 'ABSL_DEPRECATED'
#define ABSL_DEPRECATED(message) __attribute__((deprecated(message)))
                                                ^
In file included from c_src/exla/exla_client.cc:1:
In file included from c_src/exla/exla_client.h:12:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/gpu_device.h:22:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h:34:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/client/local_client.h:27:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/compiler.h:30:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_assignment.h:37:
cache/xla_extension/include/tensorflow/compiler/xla/service/memory_space_assignment.h:1000:1: warning: 'Options' defined as a struct here but previously declared as a class; this is valid, but may result in linker errors under the Microsoft C++ ABI [-Wmismatched-tags]
struct Options {
^
cache/xla_extension/include/tensorflow/compiler/xla/service/memory_space_assignment.h:37:1: note: did you mean struct here?
class Options;
^~~~~
struct
In file included from c_src/exla/exla_client.cc:1:
In file included from c_src/exla/exla_client.h:12:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/gpu_device.h:22:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h:34:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/client/local_client.h:27:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/compiler.h:30:
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_assignment.h:600:18: warning: 'set_color' is deprecated: Use Layout::memory_space instead. [-Wdeprecated-declarations]
          value->set_color(BufferValue::Color(
                 ^
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_value.h:118:3: note: 'set_color' has been explicitly marked deprecated here
  ABSL_DEPRECATED("Use Layout::memory_space instead.")
  ^
cache/xla_extension/include/absl/base/attributes.h:648:49: note: expanded from macro 'ABSL_DEPRECATED'
#define ABSL_DEPRECATED(message) __attribute__((deprecated(message)))
                                                ^
In file included from c_src/exla/exla_client.cc:1:
In file included from c_src/exla/exla_client.h:12:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/gpu_device.h:22:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h:34:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/client/local_client.h:27:
In file included from cache/xla_extension/include/tensorflow/compiler/xla/service/compiler.h:30:
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_assignment.h:603:18: warning: 'set_color' is deprecated: Use Layout::memory_space instead. [-Wdeprecated-declarations]
          value->set_color(BufferValue::Color(0));
                 ^
cache/xla_extension/include/tensorflow/compiler/xla/service/buffer_value.h:118:3: note: 'set_color' has been explicitly marked deprecated here
  ABSL_DEPRECATED("Use Layout::memory_space instead.")
  ^
cache/xla_extension/include/absl/base/attributes.h:648:49: note: expanded from macro 'ABSL_DEPRECATED'
#define ABSL_DEPRECATED(message) __attribute__((deprecated(message)))
                                                ^
5 warnings generated.
install_name_tool -change bazel-out/darwin_arm64-opt/bin/tensorflow/compiler/xla/extension/libxla_extension.so @loader_path/xla_extension/lib/libxla_extension.so -change bazel-out/darwin-opt/bin/tensorflow/compiler/xla/extension/libxla_extension.so @loader_path/xla_extension/lib/libxla_extension.so cache/libexla.so
Compiling 21 files (.ex)
Generated exla app
:ok

Section

EXLA.set_as_nx_default([:tpu, :cuda, :rocm, :host])
:host
Nx.tensor([[1, 2, 3], [4, 5, 6]])
#Nx.Tensor<
  s64[2][3]
  EXLA.Backend<host:0, 0.3682781178.2571501588.250350>
  [
    [1, 2, 3],
    [4, 5, 6]
  ]
>
Nx.tensor(10)
#Nx.Tensor<
  s64
  EXLA.Backend<host:0, 0.3682781178.2571501588.250359>
  10
>
Nx.tensor([1, 2, 3])
#Nx.Tensor<
  s64[3]
  EXLA.Backend<host:0, 0.3682781178.2571501588.250360>
  [1, 2, 3]
>
Nx.eye(3)
#Nx.Tensor<
  s64[3][3]
  EXLA.Backend<host:0, 0.3682781178.2571501583.250221>
  [
    [1, 0, 0],
    [0, 1, 0],
    [0, 0, 1]
  ]
>
Nx.iota({3, 5})
#Nx.Tensor<
  s64[3][5]
  EXLA.Backend<host:0, 0.3682781178.2571501583.250222>
  [
    [0, 1, 2, 3, 4],
    [5, 6, 7, 8, 9],
    [10, 11, 12, 13, 14]
  ]
>
a = Nx.tensor([[1, 2], [3, 4]])
b = Nx.tensor([[10, 10], [10, 10]])
c = Nx.tensor([[2, 2, 2], [3, 3, 3]])
#Nx.Tensor<
  s64[2][3]
  EXLA.Backend<host:0, 0.3682781178.2571501588.250363>
  [
    [2, 2, 2],
    [3, 3, 3]
  ]
>
Nx.add(a, b)
#Nx.Tensor<
  s64[2][2]
  EXLA.Backend<host:0, 0.3682781178.2571501583.250223>
  [
    [11, 12],
    [13, 14]
  ]
>
Nx.multiply(a, b)
#Nx.Tensor<
  s64[2][2]
  EXLA.Backend<host:0, 0.3682781178.2571501583.250224>
  [
    [10, 20],
    [30, 40]
  ]
>
Nx.dot(a, c)
#Nx.Tensor<
  s64[2][3]
  EXLA.Backend<host:0, 0.3682781178.2571501583.250225>
  [
    [8, 8, 8],
    [18, 18, 18]
  ]
>
Nx.tensor([1, 2, 3]) |> Nx.shape()
{3}
Nx.tensor([[1, 2, 3], [4, 5, 6]]) |> Nx.shape()
{2, 3}
Nx.random_uniform({2}, 0, 1, type: {:f, 32})
#Nx.Tensor<
  f32[2]
  EXLA.Backend<host:0, 0.3682781178.2571501583.250226>
  [0.5918947458267212, 0.40608131885528564]
>
Nx.random_uniform({3, 3}, 0, 100)
#Nx.Tensor<
  s64[3][3]
  EXLA.Backend<host:0, 0.3682781178.2571501583.250227>
  [
    [42, 2, 1],
    [93, 93, 48],
    [60, 41, 87]
  ]
>
Nx.random_uniform({3, 3}, 0, 100, type: {:u, 8})
defmodule MonteCarlo do
  def compute(n) when n > 0 do
    Enum.reduce(1..n, 0, fn _, acc -> acc + do_compute() end) * 4 / n
  end

  defp do_compute do
    x = :rand.uniform()
    y = :rand.uniform()
    r = x * x + y * y
    if r < 1, do: 1, else: 0
  end
end

MonteCarlo.compute(10_000_000)
3.1419408
defmodule MonteCarloNx do
  import Nx.Defn

  def compute(n) when n > 0 do
    x = Nx.random_uniform({n})
    y = Nx.random_uniform({n})

    s =
      do_compute(x, y)
      |> Nx.to_number()

    s * 4 / n
  end

  defn do_compute(x, y) do
    x = square(x)
    y = square(y)

    x
    |> Nx.add(y)
    |> Nx.less(Nx.tensor(1))
    |> Nx.sum()
  end

  defn square(vector) do
    Nx.multiply(vector, vector)
  end
end

MonteCarloNx.compute(100_000_000)
3.14184424
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment