def f(x): return reduce_sum(nonzero(x), tuple(range(len(x.shape))))
def f(x): nonzero_idx = nonzero(x) return reduce_sum(nonzero_idx)
def f(x): y = sin(x) return reduce_sum(y, axes=(0,))