How to create a matrix where each cell contain a index pair?

Questions / Help
,
1

I want to create a tensor of size m * m * 2, that simulate a square matrix where each cell contain your indexes in a pair. Something like that:

f32[3][3][2]
[
    [
        [0, 0],
        [0, 1],
        [0, 2]
    ],
    [
        [1, 0],
        [1, 1],
        [1, 2]
    ],
    [
        [2, 0],
        [2, 1],
        [2, 2]
    ]
]

I used many Nx functions (iota, broadcast, concatenate, stack, tile, reshape) trying to achieve this result, but the closest tensor I got was the following:

f32[2][3][3]
[
    [
        [0, 0, 0],
        [1, 1, 1],
        [2, 2, 2]
    ],
    [
        [0, 1, 2],
        [0, 1, 2],
        [0, 1, 2]
    ]
]

Any suggestion? Thanks

2

Why not just generating it on your own and create a tensor from it?

#!/bin/env elixir

Mix.install [:nx]

defmodule Example do
  def generic_matrix(list, acc \\ [])

  def generic_matrix([size], acc) do
    range = 0..(size - 1)

    Enum.map(range, fn index ->
      Enum.reverse([index | acc])
    end)
  end

  def generic_matrix([size | rest], acc) do
    range = 0..(size - 1)

    Enum.map(range, fn index ->
      generic(rest, [index | acc])
    end)
  end

  def square_matrix(size) do
    range = 0..(size - 1)

    Enum.map(range, fn index ->
      Enum.map(range, &[index, &1])
    end)
  end
end

3
|> tap(fn input ->
  [input, input]
  |> Example.generic_matrix()
  |> Nx.tensor()
  |> dbg()
end)
|> tap(fn input ->
  input
  |> Example.square_matrix()
  |> Nx.tensor()
  |> dbg()
end)

Notes

  1. The above code is an Elixir’s script which could be copy-pasted right into iex / livebook or saved in exs file.

  2. Using #! at the top of file you can grant executable permission using chmod and run it like ./matrix.exs.

  3. /bin/env is preferred if there is no fixed binary path i.e. unlike bash which is supposed to be in /bin directory, elixir could be placed in /usr/bin (os package manager) or $HOME/.asdf/shims (asdf version manager) or other custom path like for example when you fetch and compile sources on your own …

3

It would fit well, but I’m benchmarking so I need the best performance. I thought Nx’s built-in functions like iota/2 or broadcast/3 work faster in GPU than loops from Elixir world, but I could be wrong.

4

I guess it’s similar to implementing NIF. We need to ask ourselves how often and big data we want to transfer between CPU and GPU calculations. <1ms loss is nothing if it’s not multiplied many, many times.Square matrix for small values called just once shouldn’t be even noticeable. You could create async tasks in a loop, but I’m not sure if for such small data it’s even worth. Maybe in generic_matrix/2 version it could make sense to pattern-match checking if size is bigger than let’s say 100 and if so we could work asynchronously.

1 Like