diff --git a/src/kamae/keras/tensorflow/layers/list_max.py b/src/kamae/keras/tensorflow/layers/list_max.py index 12e91a52..161021b9 100644 --- a/src/kamae/keras/tensorflow/layers/list_max.py +++ b/src/kamae/keras/tensorflow/layers/list_max.py @@ -22,7 +22,11 @@ from kamae.keras.core.backend import TENSORFLOW_ONLY from kamae.keras.core.base import BaseLayer from kamae.keras.core.utils.input_utils import allow_single_or_multiple_tensor_input -from kamae.keras.tensorflow.utils.list_utils import get_top_n, segmented_operation +from kamae.keras.tensorflow.utils.list_utils import ( + get_top_n, + min_filter_mask, + segmented_operation, +) from kamae.keras.tensorflow.utils.transform_utils import map_fn_w_axis @@ -108,6 +112,10 @@ def compatible_dtypes(self) -> Optional[List[str]]: "float16", "float32", "float64", + "int8", + "int16", + "int32", + "int64", "string", ] @@ -145,11 +153,12 @@ def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: # Apply the mask to filter out elements less than or equal to the threshold if self.min_filter_value is not None: - mask = tf.greater_equal(val_tensor, self.min_filter_value) - neg_inf = val_tensor.dtype.min - val_tensor = tf.where(mask, val_tensor, neg_inf) - else: - val_tensor = val_tensor + mask = min_filter_mask(val_tensor, self.min_filter_value) + # The dtype minimum is only a neutral element for the reduction here: a + # real value equal to it still wins the max, so substituting it for the + # filtered entries cannot change the result. + val_tensor = tf.where(mask, val_tensor, val_tensor.dtype.min) + kept = tf.cast(mask, tf.int32) # Apply segmented calculation if self.with_segment: @@ -167,8 +176,32 @@ def _call(self, inputs: Iterable[KerasTensor], **kwargs: Any) -> KerasTensor: listwise_max = tf.broadcast_to(listwise_max, output_shape) if self.min_filter_value is not None: - fill_val = tf.constant(self.nan_fill_value, dtype=listwise_max.dtype) - listwise_max = tf.where(listwise_max != neg_inf, listwise_max, fill_val) + # Whether anything survived the filter is read from the mask rather than + # by testing the result against the sentinel. dtype.min is a legitimate + # value for the narrow integer dtypes (-128 for int8), so a sentinel test + # would silently overwrite real data with nan_fill_value. + if self.with_segment: + any_kept = map_fn_w_axis( + elems=[kept, segment_tensor], + fn=lambda x: segmented_operation(x, tf.math.unsorted_segment_max), + axis=self.axis, + fn_output_signature=tf.TensorSpec( + shape=kept.shape[self.axis :], dtype=kept.dtype + ), + ) + any_kept = tf.ensure_shape(any_kept, kept.shape) + else: + any_kept = tf.reduce_max(kept, axis=self.axis, keepdims=True) + any_kept = tf.broadcast_to(any_kept, output_shape) + + # nan_fill_value is a Python float, which tf.constant cannot convert + # directly to an integer dtype. Narrowing via numpy first handles the + # integer dtypes while preserving full precision for the float ones. + fill_val = tf.constant( + listwise_max.dtype.as_numpy_dtype(self.nan_fill_value), + dtype=listwise_max.dtype, + ) + listwise_max = tf.where(any_kept > 0, listwise_max, fill_val) return listwise_max diff --git a/src/kamae/keras/tensorflow/utils/__init__.py b/src/kamae/keras/tensorflow/utils/__init__.py index ca2949f6..880a6141 100644 --- a/src/kamae/keras/tensorflow/utils/__init__.py +++ b/src/kamae/keras/tensorflow/utils/__init__.py @@ -37,5 +37,10 @@ datetime_year, unix_timestamp_to_datetime, ) -from .list_utils import get_top_n, listify_tensors, segmented_operation # noqa: F401 +from .list_utils import ( # noqa: F401 + get_top_n, + listify_tensors, + min_filter_mask, + segmented_operation, +) from .transform_utils import map_fn_w_axis # noqa: F401 diff --git a/src/kamae/keras/tensorflow/utils/list_utils.py b/src/kamae/keras/tensorflow/utils/list_utils.py index 264fd028..9f5f48da 100644 --- a/src/kamae/keras/tensorflow/utils/list_utils.py +++ b/src/kamae/keras/tensorflow/utils/list_utils.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import math from typing import Any, Callable, List, Union import numpy as np @@ -70,6 +71,31 @@ def get_top_n( ) +def min_filter_mask(val_tensor: Tensor, min_filter_value: float) -> Tensor: + """ + Build the mask of the entries that meet a minimum value threshold. + + TensorFlow will not compare an integer tensor against a float threshold, so for + those the threshold is rounded up, which is equivalent under >= on integers. + A threshold beyond the bounds of the dtype cannot be narrowed at all, since the + cast would wrap around silently, so those two cases are answered directly. + Floats need neither adjustment because they saturate to +/-inf rather than wrap. + + :param val_tensor: Value tensor to filter. + :param min_filter_value: Minimum value an entry must meet to be kept. + :returns: Boolean tensor that is True wherever the entry is kept. + """ + if not val_tensor.dtype.is_integer: + return tf.greater_equal(val_tensor, min_filter_value) + + threshold = math.ceil(min_filter_value) + if threshold <= val_tensor.dtype.min: + return tf.ones_like(val_tensor, dtype=tf.bool) + if threshold > val_tensor.dtype.max: + return tf.zeros_like(val_tensor, dtype=tf.bool) + return tf.greater_equal(val_tensor, tf.cast(threshold, val_tensor.dtype)) + + def listify_tensors(x: Union[tf.Tensor, np.ndarray, List[Any]]) -> List[Any]: """ Converts any tensors or numpy arrays to lists for config serialization. diff --git a/src/kamae/spark/transformers/list_max.py b/src/kamae/spark/transformers/list_max.py index 2fdc2834..f2f43f2a 100644 --- a/src/kamae/spark/transformers/list_max.py +++ b/src/kamae/spark/transformers/list_max.py @@ -18,7 +18,16 @@ import tensorflow as tf from pyspark import keyword_only from pyspark.sql import DataFrame -from pyspark.sql.types import DataType, DoubleType, FloatType, StringType +from pyspark.sql.types import ( + ByteType, + DataType, + DoubleType, + FloatType, + IntegerType, + LongType, + ShortType, + StringType, +) from kamae.keras.core.backend import TENSORFLOW_ONLY from kamae.keras.tensorflow.layers import ListMaxLayer @@ -124,6 +133,10 @@ def compatible_dtypes(self) -> Optional[List[DataType]]: return [ FloatType(), DoubleType(), + ByteType(), + ShortType(), + IntegerType(), + LongType(), StringType(), ] diff --git a/tests/kamae/keras/tensorflow/layers/test_list_max.py b/tests/kamae/keras/tensorflow/layers/test_list_max.py index 7ffa8db1..6cab99a2 100644 --- a/tests/kamae/keras/tensorflow/layers/test_list_max.py +++ b/tests/kamae/keras/tensorflow/layers/test_list_max.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import numpy as np import pytest import tensorflow as tf @@ -576,6 +577,186 @@ class TestListMax: dtype=tf.float32, ), ), + # Integer values with segmentation + ( + [ + # values + tf.constant( + [ + [ + [1], + [5], + [3], + ], + [ + [7], + [2], + [9], + ], + ], + dtype=tf.int64, + ), + # segment + tf.constant( + [ + [ + [1], + [1], + [2], + ], + [ + [1], + [1], + [2], + ], + ], + dtype=tf.int64, + ), + ], + None, + None, + True, + "asc", + "int64", + "int64", + tf.constant( + [ + [ + [5], + [5], + [3], + ], + [ + [7], + [7], + [9], + ], + ], + dtype=tf.int64, + ), + ), + # Integer values with segmentation and min_filter_value, where the + # filter empties a segment entirely and nan_fill_value is applied. + ( + [ + # values + tf.constant( + [ + [ + [1], + [1], + [9], + ], + [ + [5], + [1], + [9], + ], + ], + dtype=tf.int64, + ), + # segment + tf.constant( + [ + [ + [1], + [2], + [2], + ], + [ + [1], + [2], + [2], + ], + ], + dtype=tf.int64, + ), + ], + 2, + None, + True, + "asc", + "int64", + "int64", + tf.constant( + [ + [ + [0], + [9], + [9], + ], + [ + [5], + [9], + [9], + ], + ], + dtype=tf.int64, + ), + ), + # Remaining integer widths, exercising segmentation together with + # min_filter_value emptying a segment, so that nan_fill_value is + # applied for every integer dtype the layer accepts. + *[ + ( + [ + # values + tf.constant( + [ + [ + [1], + [1], + [9], + ], + [ + [5], + [1], + [9], + ], + ], + dtype=int_dtype, + ), + # segment + tf.constant( + [ + [ + [1], + [2], + [2], + ], + [ + [1], + [2], + [2], + ], + ], + dtype=int_dtype, + ), + ], + 2, + None, + True, + "asc", + int_dtype.name, + int_dtype.name, + tf.constant( + [ + [ + [0], + [9], + [9], + ], + [ + [5], + [9], + [9], + ], + ], + dtype=int_dtype, + ), + ) + for int_dtype in [tf.int8, tf.int16, tf.int32] + ], ], ) def test_listwise_max( @@ -611,3 +792,117 @@ def test_listwise_max( output_tensor.shape == expected_output.shape ), "Output tensor shape is not the same as expected tensor shape" tf.debugging.assert_equal(output_tensor, expected_output) + + @pytest.mark.parametrize( + "dtype, nan_fill_value, expected_fill", + [ + (tf.float64, 0.1, 0.1), + (tf.float64, 123.456, 123.456), + (tf.float32, 0.1, np.float32(0.1)), + (tf.int64, 7.0, 7), + (tf.int32, 7.0, 7), + ], + ) + def test_listwise_max_nan_fill_value_dtype( + self, dtype, nan_fill_value, expected_fill + ): + """The fill value applied to a segment emptied by min_filter_value must + keep full precision on float dtypes and be usable on integer dtypes.""" + # given, a segment whose values are all removed by the filter + values = tf.constant([[[1], [1], [9]]], dtype=dtype) + segments = tf.constant([[[1], [2], [2]]], dtype=dtype) + layer = ListMaxLayer( + name="listwise_max_nan_fill_test", + min_filter_value=5, + with_segment=True, + nan_fill_value=nan_fill_value, + input_dtype=dtype.name, + output_dtype=dtype.name, + ) + # when + output_tensor = layer([values, segments]) + # then + assert ( + output_tensor.numpy().flatten()[0] == expected_fill + ), "Emptied segment was not filled with the exact nan_fill_value" + + @pytest.mark.parametrize("with_segment", [False, True]) + @pytest.mark.parametrize("dtype", [tf.int8, tf.int16, tf.int32, tf.int64]) + def test_listwise_max_dtype_minimum_is_a_valid_value(self, dtype, with_segment): + """The dtype minimum is ordinary data on the narrow integer types, -128 for + int8, so a list whose genuine maximum is that value must be returned as-is + rather than mistaken for a list emptied by the filter.""" + # given, a list whose real maximum is the dtype minimum, kept by the filter + dtype_min = dtype.min + values = tf.constant([[[dtype_min], [dtype_min], [dtype_min]]], dtype=dtype) + layer = ListMaxLayer( + name="listwise_max_dtype_min_test", + min_filter_value=dtype_min, + nan_fill_value=0.0, + with_segment=with_segment, + input_dtype=dtype.name, + output_dtype=dtype.name, + ) + inputs = values + if with_segment: + inputs = [values, tf.constant([[[1], [1], [2]]], dtype=dtype)] + # when + output_tensor = layer(inputs) + # then + tf.debugging.assert_equal( + output_tensor, + tf.constant([[[dtype_min], [dtype_min], [dtype_min]]], dtype=dtype), + ) + + @pytest.mark.parametrize( + "min_filter_value, expected", + [ + (-200.0, 3), + (-129.0, 3), + (-128.0, 3), + (127.0, 0), + (128.0, 0), + (200.0, 0), + ], + ) + def test_listwise_max_min_filter_value_outside_integer_dtype_range( + self, min_filter_value, expected + ): + """Narrowing a threshold that falls outside the value dtype wraps around, so + the out-of-range cases are decided directly: below the dtype minimum keeps + every value, above the dtype maximum keeps none of them.""" + # given, int8 values of 1, 2 and 3 against thresholds beyond the int8 bounds + values = tf.constant([[[1], [2], [3]]], dtype=tf.int8) + layer = ListMaxLayer( + name="listwise_max_out_of_range_filter_test", + min_filter_value=min_filter_value, + nan_fill_value=0.0, + input_dtype="int8", + output_dtype="int8", + ) + # when + output_tensor = layer(values) + # then + tf.debugging.assert_equal( + output_tensor, tf.constant([[[expected]] * 3], dtype=tf.int8) + ) + + @pytest.mark.parametrize("dtype", [tf.int8, tf.int16, tf.int32, tf.int64]) + def test_listwise_max_float_min_filter_value_on_integer_input(self, dtype): + """min_filter_value is declared as a float, so passing one against an + integer input column must work rather than fail the dtype comparison.""" + # given + values = tf.constant([[[1], [2], [3]]], dtype=dtype) + layer = ListMaxLayer( + name="listwise_max_float_filter_test", + min_filter_value=0.0, + nan_fill_value=-1.0, + input_dtype=dtype.name, + output_dtype=dtype.name, + ) + # when + output_tensor = layer(values) + # then + tf.debugging.assert_equal( + output_tensor, tf.constant([[[3], [3], [3]]], dtype=dtype) + ) diff --git a/tests/kamae/keras/tensorflow/test_list_utils.py b/tests/kamae/keras/tensorflow/test_list_utils.py index 71a4a06e..8636e1f1 100644 --- a/tests/kamae/keras/tensorflow/test_list_utils.py +++ b/tests/kamae/keras/tensorflow/test_list_utils.py @@ -15,7 +15,7 @@ import pytest import tensorflow as tf -from kamae.keras.tensorflow.utils import get_top_n +from kamae.keras.tensorflow.utils import get_top_n, min_filter_mask class TestGetTopN: @@ -69,3 +69,52 @@ def test_get_top_with_batch_1( ): output = get_top_n(val_tensor, axis, sort_tensor, top_n, sort_order) tf.debugging.assert_equal(output, expected_output) + + +class TestMinFilterMask: + @pytest.mark.parametrize( + "dtype, min_filter_value, expected", + [ + # Integer values cannot be compared against a float threshold directly, + # and rounding up leaves >= unchanged. + (tf.int32, 0.0, [False, True, True]), + (tf.int32, 0.5, [False, False, True]), + (tf.int32, 1.0, [False, False, True]), + (tf.int32, -0.5, [False, True, True]), + # Floats keep the fractional threshold as given. + (tf.float32, 0.5, [False, False, True]), + (tf.float64, -0.5, [False, True, True]), + ], + ) + def test_min_filter_mask_threshold(self, dtype, min_filter_value, expected): + val_tensor = tf.constant([-1, 0, 1], dtype=dtype) + output = min_filter_mask(val_tensor, min_filter_value) + tf.debugging.assert_equal(output, tf.constant(expected, dtype=tf.bool)) + + @pytest.mark.parametrize( + "min_filter_value, expected", + [ + # Narrowing these to int8 would wrap around, so they are answered without + # a cast: below the dtype minimum keeps everything, above the dtype + # maximum keeps nothing. + (-200.0, [True, True, True]), + (-129.0, [True, True, True]), + (-128.0, [True, True, True]), + (127.0, [False, False, False]), + (128.0, [False, False, False]), + (200.0, [False, False, False]), + ], + ) + def test_min_filter_mask_threshold_outside_dtype_range( + self, min_filter_value, expected + ): + val_tensor = tf.constant([1, 2, 3], dtype=tf.int8) + output = min_filter_mask(val_tensor, min_filter_value) + tf.debugging.assert_equal(output, tf.constant(expected, dtype=tf.bool)) + + def test_min_filter_mask_keeps_dtype_minimum(self): + """The dtype minimum is a real value, not a sentinel, so a threshold at the + bound must keep it.""" + val_tensor = tf.constant([-128, -127], dtype=tf.int8) + output = min_filter_mask(val_tensor, -128.0) + tf.debugging.assert_equal(output, tf.constant([True, True], dtype=tf.bool)) diff --git a/tests/kamae/spark/transformers/test_list_max.py b/tests/kamae/spark/transformers/test_list_max.py index 9f927a72..812ada45 100644 --- a/tests/kamae/spark/transformers/test_list_max.py +++ b/tests/kamae/spark/transformers/test_list_max.py @@ -664,6 +664,152 @@ def test_spark_mean_transform_with_segmentation( "double", "float", ), + # Integer values with segmentation + ( + 3, + tf.constant( + [ + [1], + [1], + [1], + [2], + [2], + [2], + ], + dtype=tf.int64, + ), + [ + # values + tf.constant( + [ + [1], + [1], + [4], + [5], + [5], + [20], + ], + dtype=tf.int64, + ), + # segment + tf.constant( + [ + [1], + [1], + [2], + [1], + [1], + [2], + ], + dtype=tf.int64, + ), + ], + None, + True, + None, + "bigint", + "bigint", + ), + # Integer values with segmentation & filter + ( + 3, + tf.constant( + [ + [1], + [1], + [1], + [2], + [2], + [2], + ], + dtype=tf.int64, + ), + [ + # values + tf.constant( + [ + [1], + [1], + [4], + [5], + [5], + [20], + ], + dtype=tf.int64, + ), + # segment + tf.constant( + [ + [1], + [1], + [2], + [1], + [1], + [2], + ], + dtype=tf.int64, + ), + ], + 2, + True, + None, + "bigint", + "bigint", + ), + # Remaining integer widths, with segmentation & filter, to check + # Spark/Keras parity for every integer dtype the transformer accepts. + *[ + ( + 3, + tf.constant( + [ + [1], + [1], + [1], + [2], + [2], + [2], + ], + dtype=tf_dtype, + ), + [ + # values + tf.constant( + [ + [1], + [1], + [4], + [5], + [5], + [20], + ], + dtype=tf_dtype, + ), + # segment + tf.constant( + [ + [1], + [1], + [2], + [1], + [1], + [2], + ], + dtype=tf_dtype, + ), + ], + 2, + True, + None, + spark_dtype, + spark_dtype, + ) + for tf_dtype, spark_dtype in [ + (tf.int8, "tinyint"), + (tf.int16, "smallint"), + (tf.int32, "int"), + ] + ], ], ) def test_list_max_transform_spark_tf_parity( @@ -734,3 +880,61 @@ def test_list_max_transform_spark_tf_parity( decimal=6, err_msg="Spark and Tensorflow transform outputs are not equal", ) + + @pytest.mark.parametrize( + "spark_dtype, tf_dtype", + [ + ("tinyint", tf.int8), + ("smallint", tf.int16), + ("int", tf.int32), + ("bigint", tf.int64), + ], + ) + def test_list_max_transform_spark_tf_parity_at_dtype_minimum( + self, spark_session, spark_dtype, tf_dtype + ): + """Spark only returns null once a filter empties a window, so a query whose + genuine maximum is the dtype minimum keeps that value. Tensorflow has to + agree rather than read the reduction result back as an emptied query.""" + # given, two queries of two items filtered at the dtype minimum, so nothing + # is removed. The first has the dtype minimum as its real maximum, the + # second keeps an ordinary value as a control. + dtype_min = tf_dtype.min + list_size = 2 + query_ids = [1, 1, 2, 2] + values = [dtype_min, dtype_min, dtype_min, 5] + transformer = ListMaxTransformer( + inputCol="input0", + outputCol="output", + inputDtype=spark_dtype, + outputDtype=spark_dtype, + queryIdCol="search_id", + minFilterValue=dtype_min, + nanFillValue=0.0, + ) + spark_df = spark_session.createDataFrame( + list(zip(query_ids, values)), ["search_id", "input0"] + ) + # when + spark_values = ( + transformer.transform(spark_df) + .select("output") + .rdd.map(lambda r: r[0]) + .collect() + ) + input_tensor = tf.reshape( + tf.constant(values, dtype=tf_dtype), (-1, list_size, 1) + ) + tensorflow_values = np.reshape( + transformer.get_keras_layer()(input_tensor).numpy(), -1 + ).tolist() + # then + assert spark_values == [dtype_min, dtype_min, 5, 5], ( + "The dtype minimum should be kept as a real value, and nanFillValue " + "should not appear when nothing is filtered" + ) + np.testing.assert_equal( + spark_values, + tensorflow_values, + err_msg="Spark and Tensorflow transform outputs are not equal", + )