keras-rs-nightly 0.0.1.dev2025022303__py3-none-any.whl → 0.0.1.dev2025022503__py3-none-any.whl

keras_rs/src/layers/modeling/dot_interaction.py

@@ -5,6 +5,7 @@ from keras import ops
 
 from keras_rs.src import types
 from keras_rs.src.api_export import keras_rs_export
+from keras_rs.src.utils.keras_utils import check_shapes_compatible
 
 
 @keras_rs_export("keras_rs.layers.DotInteraction")
@@ -44,6 +45,44 @@ class DotInteraction(keras.layers.Layer):
         self.self_interaction = self_interaction
         self.skip_gather = skip_gather
 
+    def _generate_tril_mask(
+        self, pairwise_interaction_matrix: types.Tensor
+    ) -> types.Tensor:
+        """Generates lower triangular mask."""
+
+        # If `self.self_interaction` is `True`, keep the main diagonal.
+        k = -1
+        if self.self_interaction:
+            k = 0
+
+        # Typecast k from Python int to tensor, because `ops.tril` uses
+        # `tf.cond` (which requires tensors).
+        # TODO (abheesht): Remove typecast once fix is merged in core Keras.
+        if keras.config.backend() == "tensorflow":
+            k = ops.array(k)
+        tril_mask = ops.tril(
+            ops.ones_like(pairwise_interaction_matrix, dtype=bool),
+            k=k,
+        )
+
+        return tril_mask
+
+    def _get_lower_triangular_indices(self, num_features: int) -> list[int]:
+        """Python function which generates indices to get the lower triangular
+        matrix as if it were flattened.
+        """
+        flattened_indices = []
+        for i in range(num_features):
+            k = i
+            # if `self.self_interaction` is `True`, keep the main diagonal.
+            if self.self_interaction:
+                k += 1
+            for j in range(k):
+                flattened_index = i * num_features + j
+                flattened_indices.append(flattened_index)
+
+        return flattened_indices
+
     def call(self, inputs: list[types.Tensor]) -> types.Tensor:
         """Forward pass of the dot interaction layer.
 
@@ -64,23 +103,25 @@ class DotInteraction(keras.layers.Layer):
         # Check if all feature tensors have the same shape and are of rank 2.
         shape = ops.shape(inputs[0])
         for idx, tensor in enumerate(inputs):
-            if ops.shape(tensor) != shape:
+            other_shape = ops.shape(tensor)
+
+            if len(shape) != 2:
+                raise ValueError(
+                    "All feature tensors inside `inputs` should have rank 2. "
+                    f"Received rank {len(shape)} at index {idx}."
+                )
+
+            if not check_shapes_compatible(shape, other_shape):
                 raise ValueError(
                     "All feature tensors in `inputs` should have the same "
                     f"shape. Found at least one conflict: shape = {shape} at "
-                    f"index 0 and shape = {ops.shape(tensor)} at index {idx}"
+                    f"index 0 and shape = {ops.shape(tensor)} at index {idx}."
                 )
 
-        if len(shape) != 2:
-            raise ValueError(
-                "All feature tensors inside `inputs` should have rank 2. "
-                f"Received rank {len(shape)}."
-            )
-
         # `(batch_size, num_features, feature_dim)`
         features = ops.stack(inputs, axis=1)
 
-        batch_size, _, _ = ops.shape(features)
+        batch_size, num_features, _ = ops.shape(features)
 
         # Compute the dot product to get feature interactions. The shape here is
         # `(batch_size, num_features, num_features)`.
@@ -88,28 +129,30 @@ class DotInteraction(keras.layers.Layer):
             features, ops.transpose(features, axes=(0, 2, 1))
         )
 
-        # if `self.self_interaction` is `True`, keep the main diagonal.
-        k = -1
-        if self.self_interaction:
-            k = 0
-
-        tril_mask = ops.tril(
-            ops.ones_like(pairwise_interaction_matrix, dtype=bool),
-            k=k,
-        )
-
         # Set the upper triangle entries to 0, if `self.skip_gather` is True.
         # Else, "pick" only the lower triangle entries.
         if self.skip_gather:
+            tril_mask = self._generate_tril_mask(pairwise_interaction_matrix)
+
             activations = ops.multiply(
                 pairwise_interaction_matrix,
                 ops.cast(tril_mask, dtype=pairwise_interaction_matrix.dtype),
             )
+            # Rank-2 tensor.
+            activations = ops.reshape(
+                activations, (batch_size, num_features * num_features)
+            )
         else:
-            activations = pairwise_interaction_matrix[tril_mask]
-
-        # Rank-2 tensor.
-        activations = ops.reshape(activations, (batch_size, -1))
+            flattened_indices = self._get_lower_triangular_indices(num_features)
+            pairwise_interaction_matrix_flattened = ops.reshape(
+                pairwise_interaction_matrix,
+                (batch_size, num_features * num_features),
+            )
+            activations = ops.take(
+                pairwise_interaction_matrix_flattened,
+                flattened_indices,
+                axis=-1,
+            )
 
         return activations
 

keras_rs/src/utils/keras_utils.py

@@ -2,6 +2,8 @@ from typing import Union
 
 import keras
 
+from keras_rs.src import types
+
 
 def clone_initializer(
     initializer: Union[str, keras.initializers.Initializer],
@@ -25,3 +27,18 @@ def clone_initializer(
         return initializer_class.from_config(config)
     # If we get a string or dict, just return as we cannot and should not clone.
     return initializer
+
+
+def check_shapes_compatible(
+    shape1: types.TensorShape, shape2: types.TensorShape
+) -> bool:
+    # Check rank first.
+    if len(shape1) != len(shape2):
+        return False
+
+    for d1, d2 in zip(shape1, shape2):
+        if isinstance(d1, int) and isinstance(d2, int):
+            if d1 != d2:
+                return False
+
+    return True

keras_rs/src/version.py

@@ -1,7 +1,7 @@
 from keras_rs.src.api_export import keras_rs_export
 
 # Unique source of truth for the version number.
-__version__ = "0.0.1.dev2025022303"
+__version__ = "0.0.1.dev2025022503"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.0.1.dev2025022303.dist-info → keras_rs_nightly-0.0.1.dev2025022503.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: keras-rs-nightly
-Version: 0.0.1.dev2025022303
+Version: 0.0.1.dev2025022503
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras RS team <keras-rs@google.com>
 License: Apache License 2.0

{keras_rs_nightly-0.0.1.dev2025022303.dist-info → keras_rs_nightly-0.0.1.dev2025022503.dist-info}/RECORD

@@ -4,16 +4,16 @@ keras_rs/api/layers/__init__.py,sha256=moxzBqnCCQsZiIgEUXn9mdtSALFTU5e0c4Ct8c0bW
 keras_rs/src/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/api_export.py,sha256=RsmG-DvO-cdFeAF9W6LRzms0kvtm-Yp9BAA_d-952zI,510
 keras_rs/src/types.py,sha256=UyOdgjqrqg_b58opnY8n6gTiDHKVR8z_bmEruehERBk,514
-keras_rs/src/version.py,sha256=cqw3Oj-fSCIXtwptq6VfGiNxxoW5tmFVqK53TwkAWjM,222
+keras_rs/src/version.py,sha256=MDAVWC9bHK8-Jwx-8Ya6HC2AHweZmT9LsoohVj41IQM,222
 keras_rs/src/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/modeling/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-keras_rs/src/layers/modeling/dot_interaction.py,sha256=kzxs9IUI-llDfgqy-JUr-zBez9pA7A3zIG-O0lGxwWY,5097
+keras_rs/src/layers/modeling/dot_interaction.py,sha256=Ove9USc5wFFICHfgkTfJuZjSWt7G53K3eE0HtwAO-Dw,6850
 keras_rs/src/layers/modeling/feature_cross.py,sha256=nd2SBIe1gXC-vL_BeF3vrq8XqTAb6iLxhQGmtyUl3EM,8105
 keras_rs/src/layers/retrieval/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/retrieval/brute_force_retrieval.py,sha256=mohILOt6PC6jHBztaowDbj3QBnSetuvkq55FmE39PlY,7321
 keras_rs/src/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-keras_rs/src/utils/keras_utils.py,sha256=TUEpKHWDDBmzXJ0g0cTWYwQNjsHkPQ5HMvzNshLDOSU,897
-keras_rs_nightly-0.0.1.dev2025022303.dist-info/METADATA,sha256=cd-0g1CXChmPiUzV3iKlIXZ99RKxRJvZJAEc2qYHfvs,3522
-keras_rs_nightly-0.0.1.dev2025022303.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-keras_rs_nightly-0.0.1.dev2025022303.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.0.1.dev2025022303.dist-info/RECORD,,
+keras_rs/src/utils/keras_utils.py,sha256=IjWSRieBkv7UX12qgUoI1tcOeISstCLRSTqSHpT06yE,1275
+keras_rs_nightly-0.0.1.dev2025022503.dist-info/METADATA,sha256=1XeOVbD3f-SNXAqgMsOBkZRJb_2uXbeOZY2dC2vg0ew,3522
+keras_rs_nightly-0.0.1.dev2025022503.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+keras_rs_nightly-0.0.1.dev2025022503.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
+keras_rs_nightly-0.0.1.dev2025022503.dist-info/RECORD,,