From 02c7f69fb402bc2a6add4f91d5053cf29be11d98 Mon Sep 17 00:00:00 2001
From: Matt Watson <1389937+mattdangerw@users.noreply.github.com>
Date: Sun, 16 Jul 2023 10:21:39 -0700
Subject: [PATCH] Stop tensorflow from eating all GPU memory (#473)

By default, tensorflow will consume all available GPU memory:
https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth

Say you are running on KERAS_BACKEND=jax, and jax and tf have both
been configured with GPU suport:
- keras_core will always import and initialize tensorflow
- tensorflow will use all available GPU memory
- jax will attempt any GPU allocation and immediately fail

Note this does not happen in colab because colab automatically exports
the environment variable:
TF_FORCE_GPU_ALLOW_GROWTH=true

We can do the same from keras-core.
---
 keras_core/backend/config.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/keras_core/backend/config.py b/keras_core/backend/config.py
index 5389a9638..df6c8f2a5 100644
--- a/keras_core/backend/config.py
+++ b/keras_core/backend/config.py
@@ -228,6 +228,13 @@ def standardize_data_format(data_format):
         _BACKEND = _backend
 
 
+if _BACKEND != "tensorflow":
+    # If we are not running on the tensorflow backend, we should stop tensorflow
+    # from using all available GPU memory. See
+    # https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth
+    os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
+
+
 @keras_core_export(
     [
         "keras_core.config.backend",