replace with in preprocessing tests (#787)

keras_core/distribution/distribution_lib.py

@@ -7,6 +7,7 @@
 """
 
 import contextlib
+
 import numpy as np
 
 from keras_core.backend import distribution_lib

keras_core/layers/preprocessing/category_encoding_test.py

@@ -1,5 +1,5 @@
 import numpy as np
-import tensorflow as tf
+from tensorflow import data as tf_data
 
 from keras_core import layers
 from keras_core import testing
@@ -121,7 +121,7 @@ def test_tf_data_compatibility(self):
  [0, 1, 0, 0],
  ]
  )
- ds = tf.data.Dataset.from_tensor_slices(input_data).batch(4).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_data).batch(4).map(layer)
  for output in ds.take(1):
  output = output.numpy()
  self.assertAllClose(output, expected_output)

keras_core/layers/preprocessing/discretization_test.py

@@ -1,7 +1,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -73,7 +73,7 @@ def test_tf_data_compatibility(self):
  layer = layers.Discretization(bin_boundaries=[0.0, 0.35, 0.5, 1.0])
  x = np.array([[-1.0, 0.0, 0.1, 0.2, 0.4, 0.5, 1.0, 1.2, 0.98]])
  self.assertAllClose(layer(x), np.array([[0, 1, 1, 1, 2, 3, 4, 4, 3]]))
- ds = tf.data.Dataset.from_tensor_slices(x).batch(1).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(x).batch(1).map(layer)
  for output in ds.take(1):
  output = output.numpy()
  self.assertAllClose(output, np.array([[0, 1, 1, 1, 2, 3, 4, 4, 3]]))
@@ -84,7 +84,7 @@ def test_tf_data_compatibility(self):
  np.random.random((32, 3)),
  )
  x = np.array([[0.0, 0.1, 0.3]])
- ds = tf.data.Dataset.from_tensor_slices(x).batch(1).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(x).batch(1).map(layer)
  for output in ds.take(1):
  output.numpy()
 

keras_core/layers/preprocessing/index_lookup_test.py

@@ -2,8 +2,8 @@
 
 import numpy as np
 import pytest
-import tensorflow as tf
 from absl.testing import parameterized
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -403,7 +403,7 @@ def test_saving(self):
 
  def test_adapt_with_tf_data(self):
  # Case: adapt + list inputs
- adapt_data = tf.data.Dataset.from_tensor_slices(
+ adapt_data = tf_data.Dataset.from_tensor_slices(
  ["one", "one", "one", "two", "two", "three"]
  ).batch(2)
  input_data = ["one", "two", "four"]

keras_core/layers/preprocessing/integer_lookup_test.py

@@ -1,5 +1,5 @@
 import numpy as np
-import tensorflow as tf
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -101,7 +101,7 @@ def test_tf_data_compatibility(self):
  vocabulary=[1, 2, 3, 4],
  )
  input_data = [2, 3, 4, 5]
- ds = tf.data.Dataset.from_tensor_slices(input_data).batch(4).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_data).batch(4).map(layer)
  for output in ds.take(1):
  output = output.numpy()
  self.assertAllClose(output, np.array([2, 3, 4, 0]))

keras_core/layers/preprocessing/normalization_test.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
-import tensorflow as tf
 from absl.testing import parameterized
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -64,7 +64,7 @@ def test_normalization_adapt(self, input_type):
  elif input_type == "tensor":
  data = backend.convert_to_tensor(x)
  elif input_type == "tf.data":
- data = tf.data.Dataset.from_tensor_slices(x).batch(8)
+ data = tf_data.Dataset.from_tensor_slices(x).batch(8)
 
  layer = layers.Normalization()
  layer.adapt(data)
@@ -81,7 +81,7 @@ def test_normalization_adapt(self, input_type):
  elif input_type == "tensor":
  data = backend.convert_to_tensor(x)
  elif input_type == "tf.data":
- data = tf.data.Dataset.from_tensor_slices(x).batch(8)
+ data = tf_data.Dataset.from_tensor_slices(x).batch(8)
 
  layer = layers.Normalization(axis=(1, 2))
  layer.adapt(data)

keras_core/layers/preprocessing/random_brightness_test.py

@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-import tensorflow as tf
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -55,6 +55,6 @@ def test_output(self):
  def test_tf_data_compatibility(self):
  layer = layers.RandomBrightness(factor=0.5, seed=1337)
  input_data = np.random.random((2, 8, 8, 3))
- ds = tf.data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
  for output in ds.take(1):
  output.numpy()

keras_core/layers/preprocessing/random_contrast_test.py

@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-import tensorflow as tf
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -42,6 +42,6 @@ def test_random_contrast(self):
  def test_tf_data_compatibility(self):
  layer = layers.RandomContrast(factor=0.5, seed=1337)
  input_data = np.random.random((2, 8, 8, 3))
- ds = tf.data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
  for output in ds.take(1):
  output.numpy()

keras_core/layers/preprocessing/random_crop_test.py

@@ -1,5 +1,5 @@
 import numpy as np
-import tensorflow as tf
+from tensorflow import data as tf_data
 
 from keras_core import layers
 from keras_core import testing
@@ -68,7 +68,7 @@ def test_predicting_with_longer_width(self):
  def test_tf_data_compatibility(self):
  layer = layers.RandomCrop(8, 9)
  input_data = np.random.random((2, 10, 12, 3))
- ds = tf.data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
  for output in ds.take(1):
  output = output.numpy()
  self.assertEqual(list(output.shape), [2, 8, 9, 3])

keras_core/layers/preprocessing/random_flip_test.py

@@ -1,8 +1,8 @@
 import unittest.mock
 
 import numpy as np
-import tensorflow as tf
 from absl.testing import parameterized
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -134,7 +134,7 @@ def test_tf_data_compatibility(self):
  layer = layers.RandomFlip("vertical", seed=42)
  input_data = np.array([[[2, 3, 4]], [[5, 6, 7]]])
  expected_output = np.array([[[5, 6, 7]], [[2, 3, 4]]])
- ds = tf.data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
  for output in ds.take(1):
  output = output.numpy()
  self.assertAllClose(output, expected_output)
@@ -158,7 +158,7 @@ def test_tf_data_compatibility(self):
  [[[5, 6, 7]], [[2, 3, 4]]],
  ]
  )
- ds = tf.data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
  for output in ds.take(1):
  output = output.numpy()
  self.assertAllClose(output, expected_output)

keras_core/layers/preprocessing/random_rotation_test.py

@@ -1,6 +1,6 @@
 import numpy as np
-import tensorflow as tf
 from absl.testing import parameterized
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -55,7 +55,7 @@ def test_tf_data_compatibility(self):
  input_image = np.reshape(np.arange(0, 25), (1, 5, 5, 1))
  layer = layers.RandomRotation(factor=(0.5, 0.5))
 
- ds = tf.data.Dataset.from_tensor_slices(input_image).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_image).map(layer)
  expected_output = np.asarray(
  [
  [24, 23, 22, 21, 20],

keras_core/layers/preprocessing/random_translation_test.py

@@ -1,6 +1,6 @@
 import numpy as np
-import tensorflow as tf
 from absl.testing import parameterized
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -377,6 +377,6 @@ def test_random_translation_left_numeric_constant(self, data_format):
  def test_tf_data_compatibility(self):
  layer = layers.RandomTranslation(0.2, 0.1)
  input_data = np.random.random((1, 4, 4, 3))
- ds = tf.data.Dataset.from_tensor_slices(input_data).batch(1).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_data).batch(1).map(layer)
  for output in ds.take(1):
  output.numpy()

keras_core/layers/preprocessing/random_zoom_test.py

@@ -1,6 +1,6 @@
 import numpy as np
-import tensorflow as tf
 from absl.testing import parameterized
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -95,7 +95,7 @@ def test_tf_data_compatibility(self):
  interpolation="nearest",
  fill_mode="constant",
  )
- ds = tf.data.Dataset.from_tensor_slices(input_image).batch(1).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_image).batch(1).map(layer)
  expected_output = np.asarray(
  [
  [0, 0, 0, 0, 0],

keras_core/layers/preprocessing/rescaling_test.py

@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-import tensorflow as tf
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -71,7 +71,7 @@ def test_rescaling_correctness(self):
  def test_tf_data_compatibility(self):
  layer = layers.Rescaling(scale=1.0 / 255, offset=0.5)
  x = np.random.random((3, 10, 10, 3)) * 255
- ds = tf.data.Dataset.from_tensor_slices(x).batch(3).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(x).batch(3).map(layer)
  for output in ds.take(1):
  output.numpy()
 

keras_core/layers/preprocessing/string_lookup_test.py

@@ -1,5 +1,5 @@
 import numpy as np
-import tensorflow as tf
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -54,7 +54,7 @@ def test_tf_data_compatibility(self):
  vocabulary=["a", "b", "c"],
  )
  input_data = ["b", "c", "d"]
- ds = tf.data.Dataset.from_tensor_slices(input_data).batch(3).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_data).batch(3).map(layer)
  for output in ds.take(1):
  output = output.numpy()
  self.assertAllClose(output, np.array([2, 3, 0]))

keras_core/layers/preprocessing/text_vectorization_test.py

@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-import tensorflow as tf
+from tensorflow import data as tf_data
 
 from keras_core import backend
 from keras_core import layers
@@ -71,7 +71,7 @@ def test_tf_data_compatibility(self):
  vocabulary=["baz", "bar", "foo"],
  )
  input_data = [["foo qux bar"], ["qux baz"]]
- ds = tf.data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
  for output in ds.take(1):
  output = output.numpy()
  self.assertAllClose(output, np.array([[4, 1, 3, 0], [1, 2, 0, 0]]))
@@ -83,7 +83,7 @@ def test_tf_data_compatibility(self):
  output_sequence_length=max_len,
  )
  layer.adapt(input_data)
- ds = tf.data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
+ ds = tf_data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
  for output in ds.take(1):
  output.numpy()
 
@@ -103,4 +103,4 @@ def test_tf_as_first_sequential_layer(self):
  layers.Embedding(5, 4),
  ]
  )
- model(tf.convert_to_tensor([["foo qux bar"], ["qux baz"]]))
+ model(backend.convert_to_tensor([["foo qux bar"], ["qux baz"]]))