Encoding layer should return float
PiperOrigin-RevId: 313810155
Change-Id: Id4cdd9456951f3b5a6206383a95fabf6b92ea947
diff --git a/tensorflow/python/keras/layers/preprocessing/category_encoding.py b/tensorflow/python/keras/layers/preprocessing/category_encoding.py
index b9460ed..a0b7d27 100644
--- a/tensorflow/python/keras/layers/preprocessing/category_encoding.py
+++ b/tensorflow/python/keras/layers/preprocessing/category_encoding.py
@@ -65,11 +65,11 @@
>>> layer = tf.keras.layers.experimental.preprocessing.CategoryEncoding(
... max_tokens=4)
>>> layer([[0, 1], [0, 0], [1, 2], [3, 1]])
- <tf.Tensor: shape=(4, 4), dtype=int64, numpy=
- array([[1, 1, 0, 0],
- [2, 0, 0, 0],
- [0, 1, 1, 0],
- [0, 1, 0, 1]])>
+ <tf.Tensor: shape=(4, 4), dtype=float32, numpy=
+ array([[1., 1., 0., 0.],
+ [2., 0., 0., 0.],
+ [0., 1., 1., 0.],
+ [0., 1., 0., 1.]], dtype=float32)>
Examples with weighted inputs:
@@ -286,18 +286,19 @@
binary_output = (self._output_mode == BINARY)
if self._sparse:
- return bincount_ops.sparse_bincount(
+ result = bincount_ops.sparse_bincount(
inputs,
weights=count_weights,
minlength=out_depth,
axis=-1,
binary_output=binary_output)
+ return math_ops.cast(result, K.floatx())
else:
result = bincount_ops.bincount(
inputs,
weights=count_weights,
minlength=out_depth,
- dtype=dtypes.int64,
+ dtype=K.floatx(),
axis=-1,
binary_output=binary_output)
result.set_shape(tensor_shape.TensorShape((None, out_depth)))
diff --git a/tensorflow/python/keras/layers/preprocessing/text_vectorization.py b/tensorflow/python/keras/layers/preprocessing/text_vectorization.py
index c80f998..bff7969 100644
--- a/tensorflow/python/keras/layers/preprocessing/text_vectorization.py
+++ b/tensorflow/python/keras/layers/preprocessing/text_vectorization.py
@@ -341,7 +341,7 @@
def compute_output_signature(self, input_spec):
output_shape = self.compute_output_shape(input_spec.shape.as_list())
- output_dtype = K.floatx() if self._output_mode == TFIDF else dtypes.int64
+ output_dtype = dtypes.int64 if self._output_mode == INT else K.floatx()
return tensor_spec.TensorSpec(shape=output_shape, dtype=output_dtype)
def adapt(self, data, reset_state=True):
diff --git a/tensorflow/python/keras/layers/preprocessing/text_vectorization_test.py b/tensorflow/python/keras/layers/preprocessing/text_vectorization_test.py
index 508f222..c641b2b 100644
--- a/tensorflow/python/keras/layers/preprocessing/text_vectorization_test.py
+++ b/tensorflow/python/keras/layers/preprocessing/text_vectorization_test.py
@@ -269,10 +269,10 @@
def test_layer_end_to_end_with_adapt(self, vocab_data, input_data, kwargs,
use_dataset, expected_output):
cls = get_layer_class()
- if kwargs.get("output_mode") == text_vectorization.TFIDF:
- expected_output_dtype = dtypes.float32
- else:
+ if kwargs.get("output_mode") == text_vectorization.INT:
expected_output_dtype = dtypes.int64
+ else:
+ expected_output_dtype = dtypes.float32
input_shape = input_data.shape
if use_dataset:
