mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

add parameter to initalize_cache function and added __init__ file

Browse Source
This commit is contained in:
colemanliyah 2021-07-21 22:22:40 +00:00
parent 93a20eeb93
commit 86985fd4c1
2 changed files with 19 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
jax/experimental/compilation_cache/__init__.py Normal file
View File

@ -0,0 +1,13 @@
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

9
jax/experimental/compilation_cache/compilation_cache.py
View File

@ -20,11 +20,14 @@ from typing import Optional
_cache = None
def initialize_cache(path):
"""Creates a global cache object. Should only be called once per process."""
def initialize_cache(path, max_cache_size_bytes=32 * 2**30):
"""Creates a global cache object. Should only be called once per process.
max_cache_sixe defaults to 32GiB.
"""
global _cache
assert _cache == None, f"The cache path has already been initialized to {_cache}"
_cache = FileSystemCache(path)
_cache = FileSystemCache(path, max_cache_size_bytes)
def get_executable(xla_computation, compile_options) -> Optional[xla_client.Executable]:
"""Returns the cached executable if present, or None otherwise."""