Np.argwhere operation equivalent for Nx - is this feature available currently? Is there an alternative that I’m not considering?

Nx can’t do dynamic shape modifications. You have a few options here:

  1. Crop with Evision
  2. Figure out the slice sizes outside of defn (by using Nx.argmin to figure out the starts and Nx.sum to figure out the lengths) and then use Nx.slice with those now-fixed values.