From 21656115847079981e3915f88ab4533790970f53 Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Tue, 20 Feb 2024 05:56:31 -0800 Subject: [PATCH] Fix code to populate defaults for boolean flags from environment variables. PiperOrigin-RevId: 608574620 --- jax/_src/config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/jax/_src/config.py b/jax/_src/config.py index f3af0171f929..a9d7050d23e7 100644 --- a/jax/_src/config.py +++ b/jax/_src/config.py @@ -408,6 +408,7 @@ def define_bool_state( """ if not isinstance(default, bool): raise TypeError(f"Default value must be of type bool, got {default}") + default = bool_env(name.upper(), default) name = name.lower() if upgrade: help += ' ' + UPGRADE_BOOL_HELP