Generated ipynb and md files with results for dlrm

Author: kharshith-k

examples/keras_rs/dlrm.py

@@ -1,9 +1,10 @@
 """
 Title: Ranking with Deep Learning Recommendation Model
-Author: Harshith Kulkarni
+Author: [Harshith Kulkarni](https://github.com/kharshith-k)
 Date created: 2025/06/02
-Last modified: 2025/09/01
-Description: Rank movies with DLRM using KerasRS
+Last modified: 2025/09/04
+Description: Rank movies with DLRM using KerasRS.
+Accelerator: GPU
 """
 
 """
@@ -12,13 +13,14 @@
 This tutorial demonstrates how to use the Deep Learning Recommendation Model (DLRM) to
 effectively learn the relationships between items and user preferences using a
 dot-product interaction mechanism. For more details, please refer to the
-[DLRM](https://arxiv.org/pdf/1906.00091) paper.
+[DLRM](https://arxiv.org/abs/1906.00091) paper.
 
 DLRM is designed to excel at capturing explicit, bounded-degree feature interactions and
 is particularly effective at processing both categorical and continuous (sparse/dense)
 input features. The architecture consists of three main components: dedicated input
-layers to handle diverse features, a dot-product interaction layer to explicitly model 
-feature interactions, and a Multi-Layer Perceptron (MLP) to capture implicit feature relationships.
+layers to handle diverse features (typically embedding layers for categorical features),
+a dot-product interaction layer to explicitly model feature interactions, and a
+Multi-Layer Perceptron (MLP) to capture implicit feature relationships.
 
 The dot-product interaction layer lies at the heart of DLRM, efficiently computing
 pairwise interactions between different feature embeddings. This contrasts with models
@@ -31,6 +33,7 @@
 
 ![DLRM Architecture](https://raw.githubusercontent.com/kharshith-k/keras-io/refs/heads/keras-rs-examples/examples/keras_rs/img/dlrm/dlrm_architecture.gif)
 
+
 Now that we have a foundational understanding of DLRM's architecture and key
 characteristics, let's dive into the code. We will train a DLRM on a real-world dataset
 to demonstrate its capability to learn meaningful feature interactions. Let's begin by
@@ -43,7 +46,7 @@
 
 import os
 
-os.environ["KERAS_BACKEND"] = "jax"  # `"tensorflow"`/`"torch"`
+os.environ["KERAS_BACKEND"] = "tensorflow"  # `"tensorflow"`/`"torch"`
 
 import keras
 import matplotlib.pyplot as plt
@@ -187,10 +190,7 @@ def print_stats(rmse_list, num_params, model_name):
 tutorial. Let's load the dataset, and keep only the useful columns.
 """
 
-ratings_ds = tfds.load(
-    "movielens/100k-ratings",
-    split="train"
-)
+ratings_ds = tfds.load("movielens/100k-ratings", split="train")
 
 
 def preprocess_features(x):
@@ -340,7 +340,6 @@ def __init__(
     ):
         super().__init__(**kwargs)
 
-        # Layers for categorical features (unchanged).
         self.embedding_layers = {}
         for feature_name in (
             MOVIELENS_CONFIG["categorical_int_features"]
@@ -352,8 +351,7 @@ def __init__(
                 output_dim=embedding_dim,
             )
 
-        # A single MLP for all continuous features.
-        self.continuous_mlp = keras.Sequential(
+        self.bottom_mlp = keras.Sequential(
             [
                 keras.layers.Dense(mlp_dim, activation="relu"),
                 keras.layers.Dense(embedding_dim),  # Output must match embedding_dim
@@ -362,18 +360,16 @@ def __init__(
 
         self.dot_layer = keras_rs.layers.DotInteraction()
 
-        self.dense_layers = []
+        self.top_mlp = []
         for num_units in dense_num_units_lst:
-            self.dense_layers.append(keras.layers.Dense(num_units, activation="relu"))
+            self.top_mlp.append(keras.layers.Dense(num_units, activation="relu"))
 
         self.output_layer = keras.layers.Dense(1)
 
-        # Attributes.
         self.dense_num_units_lst = dense_num_units_lst
         self.embedding_dim = embedding_dim
 
     def call(self, inputs):
-        # Process categorical features to get embeddings (unchanged).
         embeddings = []
         for feature_name in (
             MOVIELENS_CONFIG["categorical_int_features"]
@@ -394,17 +390,18 @@ def call(self, inputs):
         # Concatenate into a single tensor: (batch_size, num_continuous_features)
         concatenated_continuous = keras.ops.concatenate(continuous_inputs, axis=1)
 
-        # Pass through the single MLP to get one combined vector.
-        processed_continuous = self.continuous_mlp(concatenated_continuous)
+        # Pass through the Bottom MLP to get one combined vector.
+        processed_continuous = self.bottom_mlp(concatenated_continuous)
 
-# Combine with categorical embeddings. Note: we add a list containing the single tensor.
+        # Combine with categorical embeddings. Note: we add a list containing the
+        # single tensor.
         combined_features = embeddings + [processed_continuous]
 
         # Pass the list of features to the DotInteraction layer.
         x = self.dot_layer(combined_features)
 
-        for dense_layer in self.dense_layers:
-            x = dense_layer(x)
+        for layer in self.top_mlp:
+            x = layer(x)
 
         x = self.output_layer(x)
 
@@ -463,7 +460,7 @@ def get_dot_interaction_matrix(model, categorical_features, continuous_features)
     num_continuous_features = len(continuous_features)
     # Create a dummy input of zeros for the MLP
     dummy_continuous_input = keras.ops.zeros((1, num_continuous_features))
-    processed_continuous = model.continuous_mlp(dummy_continuous_input)
+    processed_continuous = model.bottom_mlp(dummy_continuous_input)
     all_feature_outputs.append(processed_continuous)
 
     interaction_matrix = np.zeros((num_features, num_features))