Merge branch 'main' of github.com:keras-team/keras-core

keras_core/backend/config.py

@@ -228,6 +228,13 @@ if "KERAS_BACKEND" in os.environ:
         _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",

requirements-common.txt

@@ -0,0 +1,15 @@
+namex
+black>=22
+flake8
+isort
+pytest
+pandas
+absl-py
+requests
+h5py
+protobuf
+google
+tensorboard-plugin-profile
+rich
+build
+dm-tree

requirements-cuda.txt

@@ -0,0 +1,18 @@
+# Tensorflow.
+# Cuda via pip is only on nightly right now.
+# We will pin a known working version to avoid breakages (nightly breaks often).
+tf-nightly[and-cuda]==2.14.0.dev20230712
+
+# Torch.
+# Pin the version used in colab currently (works with tf cuda version).
+--extra-index-url https://download.pytorch.org/whl/cu118
+torch==2.0.1+cu118 
+torchvision==0.15.2+cu118
+
+# Jax.
+# Pin the version used in colab currently (works with tf cuda version).
+--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+jax[cuda11_pip]==0.4.10
+
+# Common deps.
+-r requirements-common.txt

requirements.txt

@@ -1,21 +1,13 @@
+# Tensorflow.
 tensorflow
-# TODO: Use Torch CPU
-# Remove after resolving Cuda version differences with TF
+
+# Torch.
+# TODO: Use Torch CPU, remove after resolving Cuda version differences with TF
 torch>=2.0.1+cpu
 torchvision>=0.15.1
+
+# Jax.
 jax[cpu]
-namex
-black>=22
-flake8
-isort
-pytest
-pandas
-absl-py
-requests
-h5py
-protobuf
-google
-tensorboard-plugin-profile
-rich
-build
-dm-tree
+
+# Common deps.
+-r requirements-common.txt