prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  jax.config doesn't report command-line flags correctly
  Another one from @jmgilmer and I - I don't think the jax config is parsing command line flags correctly.  I don't know if this is functionally important or just a reporting error - but it is certainly important for user scripts knowing what flags have been set.
  
  If I run this script:
  ```
  from absl import app, flags
  from jax.config import config
  FLAGS = flags.FLAGS
  def main(_):
    print("FLAGS.jax_enable_x64", FLAGS.jax_enable_x64)
    print("FLAGS.jax_debug_nans", FLAGS.jax_debug_nans)
    print(config.values)
  if __name__ == "__main__":
    config.config_with_absl()
    app.run(main)
  ```
  I get the following problem: jax.config doesn't report the correct flag settings.
  ```
  > python jaxtest.py --jax_enable_x64=1 --jax_debug_nans=1
  FLAGS.jax_enable_x64 True
  FLAGS.jax_debug_nans True
  {'jax_enable_x64': 0, 'jax_xla_backend': 'xla', 'jax_backend_target': 'local', 'jax_platform_name': '', 'jax_device_values': 1, 'jax_debug_nans': 0, 'jax_disable_jit': 0}
  ```
  if I run the same with envvars instead it works:
  ```
  JAX_ENABLE_X64=1 JAX_DEBUG_NANS=1 python jaxtest.py
  FLAGS.jax_enable_x64 True
  FLAGS.jax_debug_nans True
  {'jax_enable_x64': 1, 'jax_xla_backend': 'xla', 'jax_backend_target': 'local', 'jax_platform_name': '', 'jax_device_values': 1, 'jax_debug_nans': 1, 'jax_disable_jit': 0}
  ```
  
  I've tried parsing the flags in different ways but nothing seems to fix the issue.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `jax/config.py`
  Content:
  ```
  1 # Copyright 2018 Google LLC
  2 #
  3 # Licensed under the Apache License, Version 2.0 (the "License");
  4 # you may not use this file except in compliance with the License.
  5 # You may obtain a copy of the License at
  6 #
  7 #     https://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 
  15 import sys
  16 
  17 
  18 class Config(object):
  19   def __init__(self):
  20     self.values = {}
  21     self.meta = {}
  22     self.FLAGS = NameSpace(self.read)
  23     self.use_absl = False
  24 
  25   def update(self, name, val):
  26     self.check_exists(name)
  27     if name not in self.values:
  28       raise Exception("Unrecognized config option: {}".format(name))
  29     self.values[name] = val
  30 
  31   def read(self, name):
  32     if self.use_absl:
  33       return getattr(self.absl_flags.FLAGS, name)
  34     else:
  35       self.check_exists(name)
  36       return self.values[name]
  37 
  38   def add_option(self, name, default, opt_type, meta_args, meta_kwargs):
  39     if name in self.values:
  40       raise Exception("Config option {} already defined".format(name))
  41     self.values[name] = default
  42     self.meta[name] = (opt_type, meta_args, meta_kwargs)
  43 
  44   def check_exists(self, name):
  45     if name not in self.values:
  46       raise Exception("Unrecognized config option: {}".format(name))
  47 
  48   def DEFINE_bool(self, name, default, *args, **kwargs):
  49     self.add_option(name, default, bool, args, kwargs)
  50 
  51   def DEFINE_integer(self, name, default, *args, **kwargs):
  52     self.add_option(name, default, int, args, kwargs)
  53 
  54   def DEFINE_string(self, name, default, *args, **kwargs):
  55     self.add_option(name, default, str, args, kwargs)
  56 
  57   def DEFINE_enum(self, name, default, *args, **kwargs):
  58     self.add_option(name, default, 'enum', args, kwargs)
  59 
  60   def config_with_absl(self):
  61     # Run this before calling `app.run(main)` etc
  62     import absl.flags as absl_FLAGS
  63     from absl import app, flags as absl_flags
  64 
  65     self.use_absl = True
  66     self.absl_flags = absl_flags
  67     absl_defs = { bool: absl_flags.DEFINE_bool,
  68                   int:  absl_flags.DEFINE_integer,
  69                   str:  absl_flags.DEFINE_string,
  70                   'enum': absl_flags.DEFINE_enum }
  71 
  72     for name, val in self.values.items():
  73       flag_type, meta_args, meta_kwargs = self.meta[name]
  74       absl_defs[flag_type](name, val, *meta_args, **meta_kwargs)
  75 
  76   def complete_absl_config(self, absl_flags):
  77     for name, _ in self.values.items():
  78       self.update(name, getattr(absl_flags.FLAGS, name))
  79 
  80   def parse_flags_with_absl(self):
  81     global already_configured_with_absl
  82     if not already_configured_with_absl:
  83       import absl.flags
  84       self.config_with_absl()
  85       absl.flags.FLAGS(sys.argv)
  86       already_configured_with_absl = True
  87 
  88 
  89 class NameSpace(object):
  90   def __init__(self, getter):
  91     self._getter = getter
  92 
  93   def __getattr__(self, name):
  94     return self._getter(name)
  95 
  96 
  97 config = Config()
  98 flags = config
  99 already_configured_with_absl = False
  100
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/jax/config.py b/jax/config.py
  --- a/jax/config.py
  +++ b/jax/config.py
  @@ -73,6 +73,8 @@
         flag_type, meta_args, meta_kwargs = self.meta[name]
         absl_defs[flag_type](name, val, *meta_args, **meta_kwargs)
   
  +    app.call_after_init(lambda: self.complete_absl_config(absl_flags))
  +
     def complete_absl_config(self, absl_flags):
       for name, _ in self.values.items():
         self.update(name, getattr(absl_flags.FLAGS, name))
  @@ -83,6 +85,7 @@
         import absl.flags
         self.config_with_absl()
         absl.flags.FLAGS(sys.argv)
  +      self.complete_absl_config(absl.flags)
         already_configured_with_absl = True
  
verification_info: {"golden_diff": "diff --git a/jax/config.py b/jax/config.py\n--- a/jax/config.py\n+++ b/jax/config.py\n@@ -73,6 +73,8 @@\n       flag_type, meta_args, meta_kwargs = self.meta[name]\n       absl_defs[flag_type](name, val, *meta_args, **meta_kwargs)\n \n+    app.call_after_init(lambda: self.complete_absl_config(absl_flags))\n+\n   def complete_absl_config(self, absl_flags):\n     for name, _ in self.values.items():\n       self.update(name, getattr(absl_flags.FLAGS, name))\n@@ -83,6 +85,7 @@\n       import absl.flags\n       self.config_with_absl()\n       absl.flags.FLAGS(sys.argv)\n+      self.complete_absl_config(absl.flags)\n       already_configured_with_absl = True\n", "issue": "jax.config doesn't report command-line flags correctly\nAnother one from @jmgilmer and I - I don't think the jax config is parsing command line flags correctly.  I don't know if this is functionally important or just a reporting error - but it is certainly important for user scripts knowing what flags have been set.\r\n\r\nIf I run this script:\r\n```\r\nfrom absl import app, flags\r\nfrom jax.config import config\r\nFLAGS = flags.FLAGS\r\ndef main(_):\r\n  print(\"FLAGS.jax_enable_x64\", FLAGS.jax_enable_x64)\r\n  print(\"FLAGS.jax_debug_nans\", FLAGS.jax_debug_nans)\r\n  print(config.values)\r\nif __name__ == \"__main__\":\r\n  config.config_with_absl()\r\n  app.run(main)\r\n```\r\nI get the following problem: jax.config doesn't report the correct flag settings.\r\n```\r\n> python jaxtest.py --jax_enable_x64=1 --jax_debug_nans=1\r\nFLAGS.jax_enable_x64 True\r\nFLAGS.jax_debug_nans True\r\n{'jax_enable_x64': 0, 'jax_xla_backend': 'xla', 'jax_backend_target': 'local', 'jax_platform_name': '', 'jax_device_values': 1, 'jax_debug_nans': 0, 'jax_disable_jit': 0}\r\n```\r\nif I run the same with envvars instead it works:\r\n```\r\nJAX_ENABLE_X64=1 JAX_DEBUG_NANS=1 python jaxtest.py\r\nFLAGS.jax_enable_x64 True\r\nFLAGS.jax_debug_nans True\r\n{'jax_enable_x64': 1, 'jax_xla_backend': 'xla', 'jax_backend_target': 'local', 'jax_platform_name': '', 'jax_device_values': 1, 'jax_debug_nans': 1, 'jax_disable_jit': 0}\r\n```\r\n\r\nI've tried parsing the flags in different ways but nothing seems to fix the issue.\n", "before_files": [{"content": "# Copyright 2018 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport sys\n\n\nclass Config(object):\n  def __init__(self):\n    self.values = {}\n    self.meta = {}\n    self.FLAGS = NameSpace(self.read)\n    self.use_absl = False\n\n  def update(self, name, val):\n    self.check_exists(name)\n    if name not in self.values:\n      raise Exception(\"Unrecognized config option: {}\".format(name))\n    self.values[name] = val\n\n  def read(self, name):\n    if self.use_absl:\n      return getattr(self.absl_flags.FLAGS, name)\n    else:\n      self.check_exists(name)\n      return self.values[name]\n\n  def add_option(self, name, default, opt_type, meta_args, meta_kwargs):\n    if name in self.values:\n      raise Exception(\"Config option {} already defined\".format(name))\n    self.values[name] = default\n    self.meta[name] = (opt_type, meta_args, meta_kwargs)\n\n  def check_exists(self, name):\n    if name not in self.values:\n      raise Exception(\"Unrecognized config option: {}\".format(name))\n\n  def DEFINE_bool(self, name, default, *args, **kwargs):\n    self.add_option(name, default, bool, args, kwargs)\n\n  def DEFINE_integer(self, name, default, *args, **kwargs):\n    self.add_option(name, default, int, args, kwargs)\n\n  def DEFINE_string(self, name, default, *args, **kwargs):\n    self.add_option(name, default, str, args, kwargs)\n\n  def DEFINE_enum(self, name, default, *args, **kwargs):\n    self.add_option(name, default, 'enum', args, kwargs)\n\n  def config_with_absl(self):\n    # Run this before calling `app.run(main)` etc\n    import absl.flags as absl_FLAGS\n    from absl import app, flags as absl_flags\n\n    self.use_absl = True\n    self.absl_flags = absl_flags\n    absl_defs = { bool: absl_flags.DEFINE_bool,\n                  int:  absl_flags.DEFINE_integer,\n                  str:  absl_flags.DEFINE_string,\n                  'enum': absl_flags.DEFINE_enum }\n\n    for name, val in self.values.items():\n      flag_type, meta_args, meta_kwargs = self.meta[name]\n      absl_defs[flag_type](name, val, *meta_args, **meta_kwargs)\n\n  def complete_absl_config(self, absl_flags):\n    for name, _ in self.values.items():\n      self.update(name, getattr(absl_flags.FLAGS, name))\n\n  def parse_flags_with_absl(self):\n    global already_configured_with_absl\n    if not already_configured_with_absl:\n      import absl.flags\n      self.config_with_absl()\n      absl.flags.FLAGS(sys.argv)\n      already_configured_with_absl = True\n\n\nclass NameSpace(object):\n  def __init__(self, getter):\n    self._getter = getter\n\n  def __getattr__(self, name):\n    return self._getter(name)\n\n\nconfig = Config()\nflags = config\nalready_configured_with_absl = False\n", "path": "jax/config.py"}], "after_files": [{"content": "# Copyright 2018 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport sys\n\n\nclass Config(object):\n  def __init__(self):\n    self.values = {}\n    self.meta = {}\n    self.FLAGS = NameSpace(self.read)\n    self.use_absl = False\n\n  def update(self, name, val):\n    self.check_exists(name)\n    if name not in self.values:\n      raise Exception(\"Unrecognized config option: {}\".format(name))\n    self.values[name] = val\n\n  def read(self, name):\n    if self.use_absl:\n      return getattr(self.absl_flags.FLAGS, name)\n    else:\n      self.check_exists(name)\n      return self.values[name]\n\n  def add_option(self, name, default, opt_type, meta_args, meta_kwargs):\n    if name in self.values:\n      raise Exception(\"Config option {} already defined\".format(name))\n    self.values[name] = default\n    self.meta[name] = (opt_type, meta_args, meta_kwargs)\n\n  def check_exists(self, name):\n    if name not in self.values:\n      raise Exception(\"Unrecognized config option: {}\".format(name))\n\n  def DEFINE_bool(self, name, default, *args, **kwargs):\n    self.add_option(name, default, bool, args, kwargs)\n\n  def DEFINE_integer(self, name, default, *args, **kwargs):\n    self.add_option(name, default, int, args, kwargs)\n\n  def DEFINE_string(self, name, default, *args, **kwargs):\n    self.add_option(name, default, str, args, kwargs)\n\n  def DEFINE_enum(self, name, default, *args, **kwargs):\n    self.add_option(name, default, 'enum', args, kwargs)\n\n  def config_with_absl(self):\n    # Run this before calling `app.run(main)` etc\n    import absl.flags as absl_FLAGS\n    from absl import app, flags as absl_flags\n\n    self.use_absl = True\n    self.absl_flags = absl_flags\n    absl_defs = { bool: absl_flags.DEFINE_bool,\n                  int:  absl_flags.DEFINE_integer,\n                  str:  absl_flags.DEFINE_string,\n                  'enum': absl_flags.DEFINE_enum }\n\n    for name, val in self.values.items():\n      flag_type, meta_args, meta_kwargs = self.meta[name]\n      absl_defs[flag_type](name, val, *meta_args, **meta_kwargs)\n\n    app.call_after_init(lambda: self.complete_absl_config(absl_flags))\n\n  def complete_absl_config(self, absl_flags):\n    for name, _ in self.values.items():\n      self.update(name, getattr(absl_flags.FLAGS, name))\n\n  def parse_flags_with_absl(self):\n    global already_configured_with_absl\n    if not already_configured_with_absl:\n      import absl.flags\n      self.config_with_absl()\n      absl.flags.FLAGS(sys.argv)\n      self.complete_absl_config(absl.flags)\n      already_configured_with_absl = True\n\n\nclass NameSpace(object):\n  def __init__(self, getter):\n    self._getter = getter\n\n  def __getattr__(self, name):\n    return self._getter(name)\n\n\nconfig = Config()\nflags = config\nalready_configured_with_absl = False\n", "path": "jax/config.py"}]}