Allow jax_cache_dir to be None.

Google-ML-Automation · Google-ML-Automation · commit d8c8862f9eff · 2026-02-23T12:41:41.000-08:00
PiperOrigin-RevId: 874205553
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
@@ -1301,7 +1301,7 @@ class DevelopmentAndDebugging(BaseModel):
   """General settings for development and debugging."""
 
   constant_bound_config: list = Field([], description="Legacy configuration for constant bounds.")
-  jax_cache_dir: PathStr = Field(
+  jax_cache_dir: PathStr | None = Field(
       os.path.join(os.path.expanduser("~"), "jax_cache"),
       description="Directory for JAX compilation cache.",
   )

Original file line number	Diff line number	Diff line change
`@@ -1301,7 +1301,7 @@ class DevelopmentAndDebugging(BaseModel):`
`1301`	`1301`	`"""General settings for development and debugging."""`
`1302`	`1302`
`1303`	`1303`	`constant_bound_config: list = Field([], description="Legacy configuration for constant bounds.")`
`1304`		`- jax_cache_dir: PathStr = Field(`
	`1304`	`+ jax_cache_dir: PathStr \| None = Field(`
`1305`	`1305`	`os.path.join(os.path.expanduser("~"), "jax_cache"),`
`1306`	`1306`	`description="Directory for JAX compilation cache.",`
`1307`	`1307`	`)`