Fix formatting

keras-team · Sep 19, 2023 · f53cfd1 · f53cfd1
1 parent e248ba1
commit f53cfd1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/keras_core/export/export_lib.py b/keras_core/export/export_lib.py
@@ -92,8 +92,9 @@ def __init__(self):
             )
 
         # TODO(nkovela): Make JAX version checking programatic.
-        if backend.backend()=="jax":
+        if backend.backend() == "jax":
             from jax import __version__ as jax_v
+
             if jax_v > "0.4.15":
                 raise ValueError(
                     "The export API is only compatible with JAX version 0.4.15 "

diff --git a/keras_core/export/export_lib_test.py b/keras_core/export/export_lib_test.py
@@ -1,6 +1,7 @@
 """Tests for inference-only model/layer exporting utilities."""
 import os
 import sys
+
 import numpy as np
 import pytest
 import tensorflow as tf
@@ -29,7 +30,7 @@ def get_model():
     reason="Export only currently supports the TF and JAX backends.",
 )
 @pytest.mark.skipif(
-    backend.backend()=="jax" and sys.modules["jax"].__version__ > "0.4.15",
+    backend.backend() == "jax" and sys.modules["jax"].__version__ > "0.4.15",
     reason="The export API is only compatible with JAX version <= 0.4.15.",
 )
 class ExportArchiveTest(testing.TestCase):