Being picky, but can we say `rng=rng` here, since `rng` is a keyword argument of `sample`? (see also nonlinearities.py:79)