29 Commits

Author SHA1 Message Date
Peter Hawkins
291e52a713 Fix some warnings causing CI failures on ARM.
PiperOrigin-RevId: 678454816
2024-09-24 17:25:26 -07:00
Jake VanderPlas
f090074d86 Avoid 'from jax import config' imports
In some environments this appears to import the config module rather than
the config object.
2024-04-11 13:23:27 -07:00
Jake VanderPlas
17a26235e6 Make jax.scipy.optimize test compatible with upstream scipy 2023-11-06 14:10:24 -08:00
Jake VanderPlas
d601ceb143 Test: adjust tolerances for upstream scipy.minimize 2023-09-25 12:20:34 -07:00
Jake VanderPlas
fbe4f10403 Change to simpler import for jax.config 2023-04-21 11:51:22 -07:00
Peter Hawkins
0d3277b5c3 Port more tests from jtu.cases_from_list to jtu.sample_product. 2022-10-11 21:06:08 +00:00
Peter Hawkins
c657449528 Copybara import of the project:
--
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
2022-10-10 11:35:32 -07:00
Peter Hawkins
ba557d5e1b Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.

PiperOrigin-RevId: 476167538
2022-09-22 12:27:19 -07:00
Jake VanderPlas
93deb86710 jax.scipy.optimize: fix type inconsistency 2022-06-29 09:47:53 -07:00
Peter Hawkins
a560a29e12 Increase the minimum scipy version to 1.5.
We don't have a formal support policy for scipy versions, but 1.5 dates from around the same date as the oldest supported NumPy release NEP-29 would have us support (1.20).
2022-06-24 15:07:09 -04:00
Jake VanderPlas
85660f5363 [x64] make scipy_optimize_test compatible with strict dtype promotion 2022-06-22 11:04:20 -07:00
Jake VanderPlas
3a8f478b0a [x64] make scipy_optimize_test compatible with strict dtype promotion 2022-06-21 09:29:08 -07:00
Jake VanderPlas
1567adfa47 BUG: fix incorrect order of operations in line search 2022-06-16 15:21:24 -07:00
Jeppe Klitgaard
17de89b16a feat: refactor code using pyupgrade
This PR upgrades legacy Python code to 3.7+ code using pyupgrade:
```sh
pyupgrade --py37-plus --keep-runtime-typing **.py
```

a
2022-05-17 22:14:05 +01:00
Jake VanderPlas
97512e9e44 JaxTestCase: set jax_numpy_rank_promotion='raise' by default 2022-02-14 09:22:05 -08:00
jax authors
5691010d2f Copybara import of the project:
--
d42fffd849a4bac0c0c11a3346c93f07f8c64c44 by Jake VanderPlas <jakevdp@google.com>:

JaxTestCase: set numpy_rank_promotion='raise' by default
PiperOrigin-RevId: 427896974
2022-02-10 19:08:29 -08:00
Jake VanderPlas
6324577a63 JaxTestCase: set numpy_rank_promotion='raise' by default 2022-02-10 16:54:31 -08:00
Peter Hawkins
db2e91eba2 Move jax.test_util to jax._src.test_util.
Add forwarding shims for names used by external clients of JAX in practice.

PiperOrigin-RevId: 398721725
2021-09-24 07:02:49 -07:00
Jake VanderPlas
6114e6a0d3 test_util: add decorator to set config values in test cases 2021-08-05 14:06:37 -07:00
Jake VanderPlas
30ea76cb6b disable rank promotion for jax scipy tests 2021-08-04 10:44:23 -07:00
Peter Hawkins
94446ff757 Drop Python 3.6 support.
Per the deprecation policy (https://jax.readthedocs.io/en/latest/deprecation.html),
Python 3.6 support has been due for removal since June 23, 2020.
2021-07-15 14:20:29 -04:00
Stephan Hoyer
5b3d203453
Merge branch 'master' into master 2021-06-02 01:05:47 -07:00
Stephan Hoyer
d09839124b
Fix wrong method name 2021-06-01 12:42:55 -07:00
Gregory Thornton
03a1ee9269 Update Jax linesearch to behave more like Scipy 2021-05-26 12:49:56 +01:00
Jakob Unfried
f0c7427000 add L-BFGS optimizer 2021-05-19 19:46:11 +02:00
Peter Hawkins
1de5734d63 Enforce the the args argument to jax.scipy.optimize.minimize is a tuple. 2021-02-16 15:57:46 -05:00
joshuaalbert
231168d480 all changes plus test verifcation on TPU squashed 2020-10-26 23:58:09 +01:00
Stephan Hoyer
b0ef4838d7
Fixes to test_scipy_optimize.py for Google internal tests (#3902) 2020-07-29 15:31:02 -07:00
Joshua George Albert
02009e0cf0
BFGS algorithm (#3101)
* 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>
2020-07-29 14:14:40 -07:00