4 Commits

Author SHA1 Message Date
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
Emilio Cota
6f1d82916c math_benchmark: add --set_env flag
PiperOrigin-RevId: 515417422
2023-03-09 13:04:12 -08:00
Emilio Cota
845d68b39e math_benchmark: add dot op
PiperOrigin-RevId: 515408666
2023-03-09 12:24:47 -08:00
Emilio Cota
13e875f8b8 benchmarks: add math unary benchmarks
These will be used for benchmarking FP approximations in XLA.

PiperOrigin-RevId: 503991586
2023-01-23 08:17:16 -08:00