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

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

@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.

- Create a matrix where each value is its own coordinates with
`Vix.Vips.Operation.xyz/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. - Send the resulting matrix back to
`Nx`

, or emit a binary that you can iterate over.

```
# 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

- The resulting tensor is 3d, not 2d. But that’s easily worked around.
- The coordinates start at the top left as
`0,0`

whereas Nx I believe`0,0`

is the bottom left. - 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. - 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
- 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. `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

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

Just a quick message to say “thanks”, I will definitely look into those possibilities & will share my findings here!