Rather than use torch.count_nonzero, which is non-differentiable, an alternative is to approximate this via a “differentiable relaxation using a sum of narrow Gaussian basis functions”. i.e.
def narrow_gaussian(x, ell):
    return torch.exp(-0.5 * (x / ell) ** 2)

def approx_count_nonzero(x, ell=1e-3):
    # Approximation of || x ||_0
    return len(x) - narrow_gaussian(x, ell).sum(dim=-1)

As the value for the i-th component (x[i] or xi) approaches zero, the Gaussian basis function evaluated at x tends towards 1:
limx→0 f(xi) = 1
Likewise, as the value of the i-th component deviates from zero, the Gaussian basis function evaluated at x tends towards 0:
limx→∞ f(xi) = 0
So if you have three components with values close to zero and the rest not close to zero, then the sum of all the narrow Gaussians is approximately 3, and the number of non-zero components is len(x) minus this sum.
modified from Scipy non-linear inequality constraints.ipynb as described in [FEATURE REQUEST] modify Ax API to allow for callable that evaluates a constraint and is passed to the optimizer · Issue #769 · facebook/Ax · GitHub.
Credit: David Eriksson, Senior Research Scientist at Meta