From 128961c5f508a79c3a893aa905f7e5cfc0a65b12 Mon Sep 17 00:00:00 2001
From: Haitang Hu <hthu@google.com>
Date: Fri, 29 Jul 2022 12:13:38 -0700
Subject: [PATCH] Default to mkdirs behavior in async_makedirs.

PiperOrigin-RevId: 464131053
---
 orbax/checkpoint/utils.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/orbax/checkpoint/utils.py b/orbax/checkpoint/utils.py
index ba30d328..c2f22306 100644
--- a/orbax/checkpoint/utils.py
+++ b/orbax/checkpoint/utils.py
@@ -47,8 +47,12 @@ async def run(*args, loop=None, executor=None, **kwargs):
 
 
 # TODO(cpgaffney): This functionality should be provided by an external library.
-def async_makedirs(path, *args, **kwargs):
-  return _wrap(path.mkdir)(*args, **kwargs)
+def async_makedirs(path,
+                   *args,
+                   parents: bool = True,
+                   exist_ok: bool = True,
+                   **kwargs):
+  return _wrap(path.mkdir)(*args, parents=parents, exist_ok=exist_ok, **kwargs)
 
 
 def register_ts_spec_for_serialization():