keras-rs-nightly 0.0.1.dev2025041603__py3-none-any.whl → 0.0.1.dev2025041803__py3-none-any.whl

keras_rs/src/layers/retrieval/hard_negative_mining.py

@@ -10,41 +10,6 @@ from keras_rs.src.api_export import keras_rs_export
 MAX_FLOAT = ml_dtypes.finfo("float32").max / 100.0
 
 
-def _gather_elements_along_row(
-    data: types.Tensor, column_indices: types.Tensor
-) -> types.Tensor:
-    """Gathers elements from a 2D tensor given the column indices of each row.
-
-    First, gets the flat 1D indices to gather from. Then flattens the data to 1D
-    and uses `ops.take()` to generate 1D output and finally reshapes the output
-    back to 2D.
-
-    Args:
-        data: A [N, M] 2D `Tensor`.
-        column_indices: A [N, K] 2D `Tensor` denoting for each row, the K column
-            indices to gather elements from the data `Tensor`.
-
-    Returns:
-        A [N, K] `Tensor` including output elements gathered from data `Tensor`.
-
-    Raises:
-        ValueError: if the first dimensions of data and column_indices don't
-            match.
-    """
-    num_row, num_column, *_ = ops.shape(data)
-    num_gathered = ops.shape(column_indices)[1]
-    row_indices = ops.tile(
-        ops.expand_dims(ops.arange(num_row), -1), [1, num_gathered]
-    )
-    flat_data = ops.reshape(data, [-1])
-    flat_indices = ops.reshape(
-        ops.add(ops.multiply(row_indices, num_column), column_indices), [-1]
-    )
-    return ops.reshape(
-        ops.take(flat_data, flat_indices), [num_row, num_gathered]
-    )
-
-
 @keras_rs_export("keras_rs.layers.HardNegativeMining")
 class HardNegativeMining(keras.layers.Layer):
     """Transforms logits and labels to return hard negatives.
@@ -68,21 +33,21 @@ class HardNegativeMining(keras.layers.Layer):
         negatives as well as the positive candidate.
 
         Args:
-            logits: `[batch_size, number_of_candidates]` tensor of logits.
-            labels: `[batch_size, number_of_candidates]` one-hot tensor of
-                labels.
+            logits: logits tensor, typically `[batch_size, num_candidates]` but
+                can have more dimensions or be 1D as `[num_candidates]`.
+            labels: one-hot labels tensor, must be the same shape as `logits`.
 
         Returns:
-            tuple containing:
-            - logits: `[batch_size, num_hard_negatives + 1]` tensor of logits.
-            - labels: `[batch_size, num_hard_negatives + 1]` one-hot tensor of
-                  labels.
+            tuple containing two tensors with the last dimension of
+            `num_candidates` replaced with `num_hard_negatives + 1`.
+            - logits: `[..., num_hard_negatives + 1]` tensor of logits.
+            - labels: `[..., num_hard_negatives + 1]` one-hot tensor of labels.
         """
 
         # Number of sampled logits, i.e, the number of hard negatives to be
         # sampled (k) + number of true logit (1) per query, capped by batch
         # size.
-        num_logits = ops.shape(logits)[1]
+        num_logits = ops.shape(logits)[-1]
         if isinstance(num_logits, int):
             num_sampled = min(self._num_hard_negatives + 1, num_logits)
         else:
@@ -98,14 +63,14 @@ class HardNegativeMining(keras.layers.Layer):
         # For each query, get the indices of the logits which have the highest
         # k + 1 logit values, including the highest k negative logits and one
         # true logit.
-        _, col_indices = ops.top_k(
+        _, indices = ops.top_k(
             ops.add(logits, ops.multiply(labels, MAX_FLOAT)),
             k=num_sampled,
             sorted=False,
         )
 
         # Gather sampled logits and corresponding labels.
-        logits = _gather_elements_along_row(logits, col_indices)
-        labels = _gather_elements_along_row(labels, col_indices)
+        logits = ops.take_along_axis(logits, indices, axis=-1)
+        labels = ops.take_along_axis(labels, indices, axis=-1)
 
         return logits, labels

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.dev2025041603"
+__version__ = "0.0.1.dev2025041803"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.0.1.dev2025041603.dist-info → keras_rs_nightly-0.0.1.dev2025041803.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.0.1.dev2025041603
+Version: 0.0.1.dev2025041803
 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.dev2025041603.dist-info → keras_rs_nightly-0.0.1.dev2025041803.dist-info}/RECORD

@@ -5,14 +5,14 @@ keras_rs/api/losses/__init__.py,sha256=LGW7eHQh8FbQXdMV1s9zJpbloVlz_Zlo51sorWAvF
 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=wgkFHQtzZaQah52nHJja4pYng_h75BXqZEskD1h29LI,222
+keras_rs/src/version.py,sha256=a6XGs2YLi6kf7exH9ycMMJESRdEYW2cDXyjyETsQNGQ,222
 keras_rs/src/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/feature_interaction/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/feature_interaction/dot_interaction.py,sha256=jGHcg0EiWxth6LTxG2yWgHcyx_GXrxvA61uQqpPfnDQ,6900
 keras_rs/src/layers/feature_interaction/feature_cross.py,sha256=5OCSI0vFYzJNmgkKcuHIbVv8U2q3UvS80-qZjPimDjM,8155
 keras_rs/src/layers/retrieval/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/retrieval/brute_force_retrieval.py,sha256=izdppBXxJH0KqYEg7Zsr-SL-SHgAmnFopXMPalEO3uw,5676
-keras_rs/src/layers/retrieval/hard_negative_mining.py,sha256=CY8-3W52ZBIFcEfvjXJxbFltD6ulXl4-sZCRF6stIEc,4119
+keras_rs/src/layers/retrieval/hard_negative_mining.py,sha256=IWFrbw1h9z3AUw4oUBKf5_Aud4MTHO_AKdHfoyFa5As,3031
 keras_rs/src/layers/retrieval/remove_accidental_hits.py,sha256=Z84z2YgKspKeNdc5id8lf9TAyFsbCCz3acJxiKXYipc,3324
 keras_rs/src/layers/retrieval/retrieval.py,sha256=hVOBF10SF2q_TgJdVUqztbnw5qQF-cxVRGdJbOKoL9M,4191
 keras_rs/src/layers/retrieval/sampling_probability_correction.py,sha256=80vgOPfBiF-PC0dSyqS57IcIxOxi_Q_R7eSXHn1G0yI,1437
@@ -25,7 +25,7 @@ keras_rs/src/losses/pairwise_soft_zero_one_loss.py,sha256=XBej5nybFXEQ-Vp6GLvNmq
 keras_rs/src/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/utils/keras_utils.py,sha256=IjWSRieBkv7UX12qgUoI1tcOeISstCLRSTqSHpT06yE,1275
 keras_rs/src/utils/pairwise_loss_utils.py,sha256=6eF4CTJubCySO8M5nd3_gdTlJsta_YMnwDCcdqWYGHA,3435
-keras_rs_nightly-0.0.1.dev2025041603.dist-info/METADATA,sha256=9EV3mNVpTEuZyIu5Ihha5KhjRwRJnstN4vm1iXMVvQA,3547
-keras_rs_nightly-0.0.1.dev2025041603.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-keras_rs_nightly-0.0.1.dev2025041603.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.0.1.dev2025041603.dist-info/RECORD,,
+keras_rs_nightly-0.0.1.dev2025041803.dist-info/METADATA,sha256=MQ8Edtr90TUCgLCF7VcS-D3FkP33iQmPB0R_yU5BNws,3547
+keras_rs_nightly-0.0.1.dev2025041803.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+keras_rs_nightly-0.0.1.dev2025041803.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
+keras_rs_nightly-0.0.1.dev2025041803.dist-info/RECORD,,