Perturbs the inputs of a pointwise function.

This function adds uniform noise in the range -0.5 to 0.5 to the first argument of the given function. It further replaces derivatives of the function with (analytically computed) expected derivatives w.r.t. the noise.

This is described in Sec. 4.2. in the paper

"Universally Quantized Neural Compression"
Eirikur Agustsson & Lucas Theis

f Callable. Pointwise function applied after perturbation.
x The inputs.
*args Other arguments to f.
u The noise to perturb x with. If not set and x_plus_u is not provided, it will be sampled.
x_plus_u Alternative way to provide the noise, as x+u.
expected_grads If True, will compute expected gradients.

A tuple (y, x+u) where y=f(x+u, *args) and u is uniform noise, and the gradient of y w.r.t. x uses expected derivatives w.r.t. the distribution of u.