Commits on Apr 15, 2020

1. Fixup complex values and tol in tests for jax.scipy.linalg.sparse.cg …

   The tests for CG were failing on TPUs:
   
   - `test_cg_pytree` is fixed by requiring slightly less precision than the
     unit-test default.
   - `test_cg_against_scipy` is fixed for complex values in two independent ways:
     1. We don't set both `tol=0` and `atol=0`, which made the termination
        behavior of CG (convergence or NaN) dependent on exactly how XLA handles
        arithmetic with denormals.
     2. We make use of *real valued* inner products inside `cg`, even for complex
        values. It turns that all these inner products are mathematically
        guaranteed to yield a real number anyways, so we can save some flops and
        avoid ill-defined comparisons of complex-values (see
        numpy/numpy#15981) by ignoring the complex part
        of the result from `jnp.vdot`. (Real numbers also happen to have the
        desired rounding behavior for denormals on TPUs, so this on its own would
        also fix these failures.)

   

   shoyer committed Apr 15, 2020
   Configuration menu

   • View commit details
   • Copy full SHA for 1d144b7
   • Browse repository at this point

   Copy the full SHA

   1d144b7 View commit details

   Browse the repository at this point in the history

2. comment fixup

   

   shoyer committed Apr 15, 2020
   Configuration menu

   • View commit details
   • Copy full SHA for e173f8b
   • Browse repository at this point

   Copy the full SHA

   e173f8b View commit details

   Browse the repository at this point in the history

3. fix my comment

   

   shoyer committed Apr 15, 2020
   Configuration menu

   • View commit details
   • Copy full SHA for 4559f5e
   • Browse repository at this point

   Copy the full SHA

   4559f5e View commit details

   Browse the repository at this point in the history