A recent XLA change means that certain matrix-vector products may now be lowered as matrix-matrix multiplications. This may mean that we use lower precisions where we previously did not.
PiperOrigin-RevId: 633949879
--
d39bdefb33a19e407c352df27fb04127f4fe8a1d by Peter Hawkins <phawkins@google.com>:
Migrate more tests from jtu.cases_from_list to jtu.sample_product.
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/12717 from hawkinsp:sampletest d39bdefb33a19e407c352df27fb04127f4fe8a1d
PiperOrigin-RevId: 480136538
* BFGS algorithm
Addressing https://github.com/google/jax/issues/1400
* * addresses @shoyer comments of PR
* * skip dtype checks
* * backslash in docstring
* * increase closeness tol
* * increase closeness atol to 1.6e-6
* * addresses jakevdp comments
* * same line search as scipy
* same results format
* same (and more) testing as in scipy for line search and bfgs
* 2 spacing
* documenting
* analytic hessian non default but still available
* NamedTuple classes
* * small fix in setup_method
* * small doc string addition
* * increase atol to 2e-5 for comparison
* * removed experimental analytic_hessian
* using jnp.where for all binary replace operations
* removed _nojit as this is what disable_jit does
* * fix indentation mangling
* remove remaining _nojit
* * fixing more indentation mangling
* * segregate third_party test
* * use parametrise
* * use parametrise
* * minor nitpicking
* * fix some errors
* * use _CompileAndCheck
* * replace f_0 and g_0 for (ugly) scipy variable names
* * remove unused function
* * fix spacing
* * add args argument to minimize
* adhere fmin_bfgs to scipy api
* * remove unused function
* * ignore F401
* * look into unittest
* * fix unittest error
* * delete unused function
* more adherence to scipy's api
* add scipy's old_old_fval arg though unused
* increase line_search default maxiter to 20 (10 not enough in some cases)
* * remove unused imports
* * add ord=norm to the initial convergence check
* * remove helper function
* * merge jax/master
* * Resolve a remnant conflict from merging master to solve ReadTheDocs issue.
* * Add an informative termination message and status number.
* Revert changes to unrelated files
* cleanup bfgs_minimize
* cleanup minimize.py
* Move minimize_bfgs.py to _bfgs.py
* Move more modules around
* improve docs
* high precision einsum
* Formatting in line search
* fixup
* Type checking
* fix mypy failures
* minor fixup
Co-authored-by: Stephan Hoyer <shoyer@google.com>
