TensorFlow Recommenders: two-tower retrieval with the loss wired to the metric

tfrs.Model hands the loss to a compute_loss method

This library targets people already training in TensorFlow who want a two-tower retrieval model without hand-writing the scoring contract. Its entry point is a subclass of tfrs.Model rather than a Keras layer, and the whole contract is one method:

class Model(tfrs.Model):

  def __init__(self):
    super().__init__()

    # Set up user representation.
    self.user_model = tf.keras.layers.Embedding(
        input_dim=2000, output_dim=64)

That method returns the loss itself, so the compile call passes an optimizer and nothing else. No loss name, no metrics list:

model = Model()
model.compile(optimizer=tf.keras.optimizers.Adagrad(0.5))

Keras asks the subclass what to optimise, and the same object carries the evaluation path later. Readers who have written a custom train_step will recognise the inversion immediately. Readers who have only ever named a loss string in compile will need a minute with the quick start, because the direction of control runs the other way.

The candidate list is embedded inside the task constructor

A retrieval metric is assembled in the same constructor, and the candidate set is part of that call:

    self.task = tfrs.tasks.Retrieval(
        metrics=tfrs.metrics.FactorizedTopK(
            candidates=movies.batch(128).map(self.item_model)
        )
    )

The map runs while the task is being created, which is before compile and before the first batch, so the item side of the metric is whatever item_model returns at that moment. Catalogue tensors stay resident for the whole run, since every scored user is compared against all of them, in batches of 128.

Comfortable for a Movielens 100K example, less so for a catalogue in the millions, where the metric wants more memory than the model itself. Deployment is named in the workflow the README describes, yet nothing there explains how to shrink the candidate set at serving time. Nothing in the quick start answers that question, so read the in-memory metric as an evaluation aid for small catalogues rather than as a plan for production traffic.

Two embedding tables fixed at 2000 rows and 64 columns

The item side repeats the same literal sizes:

    # Set up movie representation.
    self.item_model = tf.keras.layers.Embedding(
        input_dim=2000, output_dim=64)

Nothing in the example derives that width from the data. Ids arrive as integers because the datasets are mapped through tf.strings.to_number before the model sees them, and the table size stays a number typed into the class. There is no vocabulary layer here, no count of distinct user ids, no lookup table, and no guard of any kind. These two tables are the entire model as written.

Any id at or above 2000 simply has no row in either table, and the example contains nothing that would report it. Point the same class at a real catalogue and the first edit you owe somebody is replacing those literals with a vocabulary built from your own data. Leaving them in place is how a demo of 2000 users turns into a model that cannot represent the user who signs up next.

The train and test halves share the same users

The split happens on rows, not on people:

# Randomly shuffle data and split between train and test.
tf.random.set_seed(42)
shuffled = ratings.shuffle(100_000, seed=42, reshuffle_each_iteration=False)

train = shuffled.take(80_000)
test = shuffled.skip(80_000).take(20_000)

Five epochs of fitting later, the script evaluates the held-out rows:

# Train.
model.fit(train.batch(4096), epochs=5)

# Evaluate.
model.evaluate(test.batch(4096), return_dict=True)

No grouping by user happens before the cut, so a viewer who rated forty films contributes training rows and test rows at once, and FactorizedTopK scores that viewer against a catalogue where the model has already met them. Metrics printed by that script therefore describe warm users, and the README does not offer a user-level or time-based alternative. With the seed fixed and reshuffle_each_iteration off, the split is at least repeatable, which is the one thing in its favour. A repeatable leak is still a leak, and closing it costs a handful of lines the example never shows.

One pip line, but the dependency floor is 2.9.0

Installation is a single command:

pip install tensorflow-recommenders

The README frames it as: make sure you have TensorFlow 2.x installed, then install from pip. The dependency file is narrower than that sentence:

tensorflow >= 2.9.0; sys_platform != 'darwin'
tensorflow-macos >= 2.9.0; sys_platform == 'darwin'
absl-py >= 0.1.6
tf-keras

Three things follow. The floor is 2.9.0, so an environment sitting on an older 2.x has to move before the resolver is happy. On macOS the named distribution becomes tensorflow-macos, a separate package name on the same import path, and the README sentence about TensorFlow 2.x never mentions that branch. Then tf-keras sits in the base list with no extras marker, no comment and no line of prose anywhere in the README explaining what pulls it in or what changes if you remove it. The quick start imports tensorflow, tensorflow_datasets and tensorflow_recommenders, and nothing documents that chain end to end.

Four packages sit behind the docs extra in setup.py

Packaging is a plain setup.py, and two of its details shape what a default install gives you:

    extras_require={
        "docs": ["fire", "annoy", "scann == 1.2.*", "tensorflow-ranking"],
    },

Everything a plain install pulls is the four lines in the dependency file, so annoy and scann == 1.2.*, the approximate search backends, arrive only when you ask for that extra, alongside fire and tensorflow-ranking. Asking for a docs extra and receiving two search indexes and a separate ranking library is a strange bundle, and the scann pin is written as 1.2.* with no lower or upper neighbour to guide a later change.

The requirement list is not a literal either:

    install_requires=pathlib.Path("requirements.txt").read_text().splitlines(),

so the constraints live in a text file read at build time rather than inside the package metadata, and every reviewer who wants to know what an install pulls has to open a second file. setup.py declares VERSION = 0.7.7 and the classifier list marks the project as Development Status :: 3 - Alpha, which is the project's own statement about how finished it considers itself. Nothing in the repository is marked archived, and the last push landed on 2026-10-06. Terms are Apache-2.0, and the classifier set repeats it as License :: OSI Approved :: Apache Software License, so a product that ships a model trained with this library carries no extra obligation to publish anything.

The deployment claim outruns the repository layout

The README says the library helps with the full workflow, naming deployment as one of its stages. Top-level entries answer for several of the others: a docs directory, a tools directory, a requirements file, a setup.py, a .pylintrc, a citation file, an assets directory and the package itself. There is no serving module at the top level, no SavedModel export in the quick start, and no port, no model server configuration and no conversion command anywhere in the README.

That gap matters more here than in most libraries, because retrieval is two systems rather than one: a model that ranks, and a service that answers. The example stops at model.evaluate, which prints metrics into a Python process that then exits. Anyone wiring this into a serving stack learns the export path from the tutorials at tensorflow.org, and the API reference link in the README points at the same external site rather than at a file in the repository. Budget for reading docs/ before you commit a service to this code.

Exact dot products in memory, or an approximate index

Two approaches sit beside this one, and picking wrongly costs more than learning the API.

The first is an approximate nearest neighbour index. Packaging names annoy and scann == 1.2.* in the docs extra, so the intent to support that path is declared, and neither library arrives with a default install. The approach differs from the quick start in exactly one way: instead of scoring every catalogue row against every user vector, you build an index over item vectors and ask it for neighbours, which trades exactness for a search that stays flat as the catalogue grows. Neither the README nor the extra says which parts of the library the index drives, so that route means reading docs/ before you write production code against it.

The second is plain Keras. Count what the quick start adds around its two embedding layers and an optimizer: a base class, one method, and a task object that couples a loss to a metric. If your goal is a two-tower model with your own evaluation loop and an export path your platform already understands, that variant is code you own end to end, with no candidate catalogue pinned inside a metric constructor and no 2000-row table to remember.

Editorial conclusion

Adopt TensorFlow Recommenders when the model you actually need is two-tower retrieval trained inside TensorFlow and you want the loss and the top-K metric bound to one task object. Leave it when ranking, multi-stage retrieval or serving carry the real weight, since the quick start stops at model.evaluate and the repository ships no serving file. Before installing, confirm that requirements.txt resolves for your platform, then rebuild both 2000-row embedding tables from your own id vocabulary and split the ratings by user before trusting a single number.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. README
4. Releases
5. tensorflow/recommenders on GitHub