Skip to content

Instantly share code, notes, and snippets.

@hugobarauna
Last active April 22, 2024 17:36
Show Gist options
  • Save hugobarauna/1f62ed805e03ab839c83819d674b660d to your computer and use it in GitHub Desktop.
Save hugobarauna/1f62ed805e03ab839c83819d674b660d to your computer and use it in GitHub Desktop.
A custom Kino to visualize a 2-dimensional Nx tensor as a matrix inside Livebook

Custom Kino to visualize a 2-dimensional Nx tensor as a matrix

Mix.install(
  [
    {:nx, "~> 0.7.1"},
    {:kino, "~> 0.12.3"}
  ],
  consolidate_protocols: false
)

How it works

Livebook uses $\TeX$ syntax for math inside your Markdown cells.

To write your own, put your math expressions between $$ if you want display math.

That's the syntax to represent matrices:

$$
\begin{bmatrix}
a & b \\
c & d
\end{bmatrix}
$$

You can explore all supported expressions in the KaTeX documentation.

We can leverate that to visualize 2-dimensionals Nx tensors as matrices

defimpl Kino.Render, for: Nx.Tensor do
  def to_livebook(tensor) do
    if tuple_size(Nx.shape(tensor)) == 2 do
      latex_matrix =
        Nx.to_list(tensor)
        |> Enum.map(fn row ->
          Enum.join(row, " & ")
        end)
        |> Enum.join(~S[\\])
        |> then(fn matrix ->
          ~S[$$\begin{bmatrix}] <> matrix <> ~S[\end{bmatrix}$$]
        end)

      Kino.Layout.tabs(Matrix: Kino.Markdown.new(latex_matrix), raw: Kino.Inspect.new(tensor))
      |> Kino.Render.to_livebook()
    else
      Kino.Inspect.new(tensor)
      |> Kino.Render.to_livebook()
    end
  end
end
Nx.tensor([[1, 2], [3, 4]])
Nx.tensor([[[1, 2], [3, 4], [5, 6]], [[-1, -2], [-3, -4], [-5, -6]]])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment