Filtering - I would like to remove zero values from my tensors

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!

I like your suggestioon, but I still have a divide by zero problem in my implementation. Here is my trivial test case. Should this be working?

def calc_function(distances) do
    Nx.select(
      Nx.not_equal(distances, 0),
      Nx.divide(1, distances),
      0
    )
end

This is my actual function to evaluate on true:

Nx.cos(Nx.multiply(6,distances)) |> Nx.divide(distances)

@seanmor5
It looks like I may not have replied to your post correctly with the results my test implementation below and you may have missed it so I am trying again. I apologize if you did see my response but just have not had a chance to get back to me!