A custom Kino to visualize a 2-dimensional Nx tensor as a matrix inside Livebook

nx_tensor_as_latex_table.livemd

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]]])