Peter Hawkins 76cda0ae07 Update flags to use the ABSL typed flag API. ...

Change flags to use the newer definition style where the flag is read via a typed FlagHolder object returned by the DEFINE_... function. The advantage of doing this is that `flag.value` has a type known to the type checker, rather than reading it as an attr out of a gigantic config dictionary.

For jax.config flags, define a typed FlagHolder object that is returned when defining a flag, matching the ABSL API.

Move a number of flags into the file that consumes them. There's no reason we're defining every flag in `config.py`.

This PR does not change the similar "state" objects in `jax.config`. Changing those is for a future PR.

PiperOrigin-RevId: 551604974

2023-07-27 12:15:58 -07:00

..

api_benchmark.py

Improve the shape incompatible error message by adding the argument/result name path to it.

2023-05-04 21:50:04 -07:00

math_benchmark.py

Update flags to use the ABSL typed flag API.

2023-07-27 12:15:58 -07:00