Parametrize the number of random bits used in stochastic rounding

[awf]

Small change to allow the number of random bits used in stochastic rounding to be specified.
The calculation is slightly modified to avoid bias, e.g. the following graph shows how the stochastically rounded values approach the full-precision values:

from matplotlib import pyplot as plt
import torch
import unit_scaling
from unit_scaling.formats import FPFormat

srnumbits = 2

fmt = FPFormat(5,2, "stochastic", srnumbits)

vs = torch.arange(0.9, 2.7, 1.0 / 117, dtype=torch.float32)

# Stochastic rounding
rvs = torch.zeros_like(vs)
nsamples = 5000
for n in range(nsamples):
    rvs += fmt.quantise(vs)
rvs /= nsamples

# round to nearest
fmt_nearest = FPFormat(5,2, "nearest")
rn = fmt_nearest.quantise(vs)

plt.plot(vs, vs, label="True value", color="k")
plt.plot(vs, rn, label="Round to nearest", color="orange")
plt.plot(vs, rvs, label=f"Stochastic, {srnumbits} bits", color="green")
plt.xlabel("Input float value")
plt.ylabel(f"Average of rounded values, {nsamples} samples")
plt.legend()

[thecharlieblake]

Thanks for this Andrew. Keen to find a case where we need stochastic rounding and play around with how many bits