Add experimental precision doubling transform

[jakevdp]

This PR adds an experimental precision doubling transform, following the basic approach outlined in Dekker 1971 (pdf). When this transform is applied, the number of significant bits is approximately doubled compared to the base operation.

Simple demo:

In [1]: import jax.numpy as jnp                                                                                        

In [2]: from jax.experimental.doubledouble import doubledouble                                                         

In [3]: def f(a, b): 
   ...:     return a + b - a 
   ...:                                                                                                                

In [4]: f(1E20, 1.0)  # float64 loses precision
Out[4]: 0.0

In [5]: g = doubledouble(f)(1E20, 1.0)
Out[5]: DeviceArray(1., dtype=float64)

This initial experiment supports basic arithmetic operators and inequalities.

[shoyer]

Very cool!

[shoyer]

You can double-double, but can you double-double-double? Might be a nice consistency test, and conceivably if somebody cares about double-double they might be interested in even higher precision, too.

[mattjj]

You can double-double, but can you double-double-double? Might be a nice consistency test, and conceivably if somebody cares about double-double they might be interested in even higher precision, too.

Yes! That was a key goal (and win) from doing this in JAX.

I'm sure it can work, though I'm not sure if there are any fiddly things left to get right in the current implementation (as of yesterday afternoon it was a WIP).

[jakevdp]

One main question I have: I believe I'll need to flatten inputs to make this as general as possible – any pointers on that?

[mattjj]

@jakevdp the general pattern is in most api.py transformations, and might look something like this (untested):

def doubledouble(f):
  def wrapped(*args):
    args_flat, in_tree = tree_flatten(args)
    f_flat, out_tree = flatten_fun_nokwargs(lu.wrap_init(f), in_tree)
    arg_pairs = [(x, jnp.zeros_like(x)) for x in args_flat]
    out_pairs_flat = doubling_transform(f_flat).call_wrapped(*arg_pairs)
    out_flat = [hi + lo for hi, lo in out_pairs_flat]
    out = tree_unflatten(out_tree(), out_flat)
    return out
  return wrapped

[mattjj]

The logic there is: we basically want to generate a function f_flat that does a bit of processing to its input before calling f and processing to its output before returning, something like

def f_flat(*flat_args):
  args = unflatten(flat_args)
  out = f(*args)
  out_flat = flatten(out)
  return out_flat

What should we flatten to and from? Well, the caller needs to flatten the args it receives, so it can pass them to f_flat (which then turns around and unflattens them to call f). That's why we pass in in_tree from the caller: f_flat needs to know how to unflatten its flat inputs. The output side is similar but trickier: we need to plumb out the out_tree somehow, smuggling it out of the function so the caller can unflatten the flat outputs we give it.

That's what flatten_fun_nokwargs in api_util.py does:

@lu.transformation_with_aux
def flatten_fun_nokwargs(in_tree, *args_flat):
  py_args = tree_unflatten(in_tree, args_flat)
  ans = yield py_args, {}
  out_flat, out_tree = tree_flatten(ans)
  yield out_flat, out_tree

It's written in an unfamiliar way, but the logic is just the above: we pass in in_tree and args_flat, unflatten to get args, call the underlying function f, then flatten its outputs and return them while smuggling out out_tree. The first yield statement is about calling the underlying function, while the second yield statement is returning to the caller.

[mattjj]

This is so cool! Nice work, I can't wait to learn more. I also can't wait for doubledouble(doubledouble(f)), but no need to block on it :)

[jakevdp]

Tests pass internally; I'm going to merge and then keep iterating.

[gugar20]

Is it possible to use this transformation on linear algebra functions? I would like to apply it to matrix inversion.

[shoyer]

In theory, yes. In practice, no: doubling precision rules have not yet been defined for any linear algebra routines.

…

On Sat, Apr 17, 2021 at 12:48 PM gugar20 ***@***.***> wrote: Is it possible to use this transformation on linear algebra functions? I would like to apply it to matrix inversion. — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#3465 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJJFVVDP667YZYEHMOV3XDTJHQYZANCNFSM4N775S4Q> .