This is not possible as it would require runtime dynamic shapes—-which Nx does not support right now. My guess is though you don’t need to filter out nonzero values but a way to ignore them in a downstream calculation. Usually you can accomplish this with Nx.select. What exactly are you trying to do?

Thank you, @seanmor5
“…a way to ignore them in a downstream calculation.” This is correct. I take a function of each of the distance values and then sum across each row. However, the function has the distance in the denominator, which is not good if the distance is zero.

For purposes of discussion, the function could be Nx.divide(1,d) where d is an individual value from the tensor. The diagonal will always be zeros, but it is possible they could show up in other regions as well.

This is perfect for a Nx.select then. We do this a lot in Axon and other apps:

Nx.select(
  Nx.not_equal(x, 0),
  Nx.divide(x, d),
  0
)

Note that you can replace 0 with whatever value you want. It’s important that you’re clever with choosing the negative case. Without more context it’s hard to say what it should be, but if the result of the divide operation is immediately summed then 0 values will be ignored. Depending on the use case you might also choose to make the masked case a small fuzz value like 1.0e-15, something finite but really large like 1e100 or -1e100, or something else. Again it all depends on context!