How to list all (i,j) coordinates for cells equal to 1?

Elixir Framework Forums Nx Forum
,
1

Hello!

I have a largish 2D tensor (15k x 15k) ultimately filled with 0 and 1.

I would like to stream, across the 225M values, the coordinates (i,j) of all the cells filled with 1.

What is the best way to do that with Nx?

Thanks!

– Thibaut

2

@thbar, I’m not in any way proficient in Nx, or matrix math. But I like a puzzle. Here is one (incomplete, probably not scalable for your needs) solution using Image. Image uses Vix that in is a NIF interface to the fabulous libvips. Which is a long way to say that libvips has a function that generates a matrix where its values are their own coordinates. And since it’s a NIF, and it can share tensors with Nx in a zero-copy way, it’s quite efficient.

Summary of partial solution

  1. Create a matrix where each value is its own coordinates with Vix.Vips.Operation.xyz/2.
  2. Multiply that matrix with your mask matrix. All the 0 mask values will zero out the coordinates and the 1 mask value will keep them in place.
  3. Send the resulting matrix back to Nx, or emit a binary that you can iterate over.

Implementation of partial solution

# Create a mask matrix for the purposes of the example
# The diagonal is set to 1
iex> mask = Nx.tensor([[1,0,0],[0,1,0],[0,0,1]], type: :u8)

# Reshape it to that its a 3d matrix that VIx/Image can use
iex> three_d_mask = Nx.reshape mask, {3, 3, 1}

# Create the coordinate matrix.
# The value of each coordinate is its own coordinate
iex> {:ok, xyz} = Vix.Vips.Operation.xyz(3, 3)

# Multiply the matrices together
iex> {:ok, multiplied} = Vix.Vips.Operation.multiply xyz, three_d_mask

# Send the matrix to Nx
# This is a zero-copy reference only
iex> Image.to_nx(multiplied)
{:ok,
 #Nx.Tensor<
   u32[height: 3][width: 3][bands: 2]
   EXLA.Backend<host:0, 0.854855650.1658716180.76475>
   [
     [
       [0, 0],
       [0, 0],
       [0, 0]
     ],
     [
       [0, 0],
       [1, 1],
       [0, 0]
     ],
     [
       [0, 0],
       [0, 0],
       [2, 2]
     ]
   ]
 >}

Now you can see that all the non-zero coordinates remain as per the mask

Issues with the example

  1. The resulting tensor is 3d, not 2d. But that’s easily worked around.
  2. The coordinates start at the top left as 0,0 whereas Nx I believe 0,0 is the bottom left.
  3. If the 0,0 coordinate is selected (has a mask value of 1) then it won’t show up because 0 * 1 = 0 so it would need special handling. Or more likely a more complex operation than just multiplication.
  4. I suspect this will be a struggle on a 15kx15k tensor because you need two of them (the coordinate matrix and the mask matrix) and then the result matrix
  5. The calculations are CPU-based, not GPU-based. libvips does use highway to provide some SIMD performance improvements but it won’t match a GPU-optimised implementation.
  6. libvips is primarily a demand-drive stream-oriented library so it could be great for a streaming calculator. For the moment, the implementation in Image can only stream image formats, not raw pixel data. If that changes, then the coordinate matrix and the mask matrix can be streamed into the resulting matrix which would save a lot of memory.

I hope that is some very small way this helps get you started.

4 Likes
3

The fabulous @akash-akya has submitted this PR to libvips to stream raw pixels from a transformation pipeline. Which means when accepted, there will be a genuine streaming answer to the original question. Very very cool.

2 Likes