Phoenix Retrieval Forward

Source: phoenix/recsys_retrieval_model.py#L38-L372 (plus recsys_model reducers)
Symbols: B = batch, S = history_seq_len, C = candidate_seq_len, D = emb_size, A = num_actions, Vps = product_surface_vocab_size, H = item_hashes + author_hashes, N = corpus size, K = top_k.
Two-tower split: User tower builds `user_representation [B, D]` via reducer blocks + Transformer; Candidate tower builds normalized item vectors `[B, C, D]` (offline corpus path).
Retrieval score: `scores = user_representation @ corpus_embeddings^T` with optional `corpus_mask`, then `jax.lax.top_k(scores, K)`.
activation tensor
mask / boolean tensor
parameter tensor
scalar / constant
named tuple / structured object

1) Entry Inputs and Constants

L316 · batch fields used in retrieval path
fields
5
user_hashes [B, U]
history_post_hashes [B, S, Hitem]
history_product_surface [B, S]
history_actions [B, S, A]
candidate_post_hashes [B, C, Hitem]
minimal fields referenced by build_user_representation / build_candidate_representation
L317 · recsys_embeddings fields used in retrieval path
fields
5
user_embeddings [B, U, D]
history_post_embeddings [B, S, Hitem, D]
history_author_embeddings [B, S, Hauthor, D]
candidate_post_embeddings [B, C, Hitem, D]
candidate_author_embeddings [B, C, Hauthor, D]
embedding tensors pre-looked-up from hash tables
L318 · corpus_embeddings
N
D
[N, D] normalized corpus vectors
L319 · top_k
int scalar (K)
L320 · corpus_mask (optional)
N
1
[N] optional boolean mask
L34-L35 · EPS, INF
scalar constants (1e-12, 1e12)

2) User-Tower Preprocessing (History Features)

L195-L200 · product_surface_embedding_table = hk.get_parameter("product_surface_embedding_table", [Vps, D])
Vps
D
[Vps, D]
L202 · history_product_surface_one_hot = one_hot(batch.history_product_surface, Vps)
S
Vps
[B, S, Vps]
L203-L204 · history_product_surface_embeddings = dot(history_product_surface_one_hot, product_surface_embedding_table)
S
D
[B, S, D]
L171-L176 · history_actions_projection = hk.get_parameter("action_projection", [A, D])
A
D
[A, D]
L178 · history_actions_signed = (2 * history_actions - 1)
S
A
[B, S, A]
L179 · history_actions_embeddings (pre-mask) = dot(history_actions_signed, history_actions_projection)
S
D
[B, S, D]
L181 · history_actions_valid_mask = any(history_actions, axis=-1, keepdims=True)
S
1
[B, S, 1]
L182-L184 · history_actions_embeddings = history_actions_embeddings (pre-mask) * history_actions_valid_mask
S
D
[B, S, D]

3) User Tower: build_user_representation

recsys_model:L79-L112 · block_user_reduce -> user_embeddings
1
D
[B, 1, D]
recsys_model:L79-L112 · block_user_reduce -> user_padding_mask
1
1
[B, 1]
recsys_model:L122-L171 · block_history_reduce -> history_embeddings
S
D
[B, S, D]
recsys_model:L122-L171 · block_history_reduce -> history_padding_mask
S
1
[B, S]
L255 · embeddings = concat([user_embeddings, history_embeddings], axis=1)
T=1+S
D
[B, 1+S, D]
L256 · padding_mask = concat([user_padding_mask, history_padding_mask], axis=1)
T=1+S
1
[B, 1+S]
L258-L262 · model_output = self.model(embeddings, padding_mask, candidate_start_offset=None)
T=1+S
D
TransformerOutput with embeddings [B, 1+S, D]
L264 · user_outputs = model_output.embeddings
T=1+S
D
[B, 1+S, D]
L266 · mask_float = padding_mask.astype(float32)[:, :, None]
T=1+S
1
[B, 1+S, 1]
L267 · user_embeddings_masked = user_outputs * mask_float
T=1+S
D
[B, 1+S, D]
L268 · user_embedding_sum = sum(user_embeddings_masked, axis=1)
1
D
[B, D]
L269 · mask_sum = sum(mask_float, axis=1)
1
1
[B, 1]
L270 · user_representation_pre_norm = user_embedding_sum / maximum(mask_sum, 1.0)
1
D
[B, D]
L272 · user_norm_sq
1
1
[B, 1]
L273 · user_norm = sqrt(maximum(user_norm_sq, EPS))
1
1
[B, 1]
L274 · user_representation = user_representation_pre_norm / user_norm
1
D
[B, D] L2-normalized
L276 · return (user_representation, user_norm)
1
D
1
1
.
[B, D]
.
[B, 1]

4) Candidate Tower (Offline Corpus Build Path)

L298 · candidate_post_embeddings
C
Hitem·D
[B, C, Hitem, D]
L299 · candidate_author_embeddings
C
Hauthor·D
[B, C, Hauthor, D]
L301-L303 · post_author_embedding = concat(axis=2)
C
H·D
[B, C, H, D], H = Hitem + Hauthor
L308 call site: CandidateTower.__call__(post_author_embedding) (body in section 5)
Section 4 outputs from build_candidate_representation
L308 · output candidate_representation
C
D
[B, C, D] L2-normalized
L310 · output candidate_padding_mask
C
1
[B, C] boolean

5) Candidate Tower Internals (CandidateTower.__call__)

Input to CandidateTower.__call__
L301-L303 · input post_author_embedding
C
H·D
[B, C, H, D], H = Hitem + Hauthor
CandidateTower.__call__ body (L68-L99)
L70 · reshape(post_author_embedding)
C
HD
[B, C, H*D] (primary 4D branch)
L77-L82 · proj_1
HD
2D
[H*D, 2D]
L84-L89 · proj_2
2D
D
[2D, D]
L91 · hidden_linear = dot(input, proj_1)
C
2D
[B, C, 2D]
L92 · hidden = silu(hidden_linear)
C
2D
[B, C, 2D]
L93 · candidate_embeddings = dot(hidden, proj_2)
C
D
[B, C, D]
L95 · candidate_norm_sq
C
1
[B, C, 1]
L96 · candidate_norm
C
1
[B, C, 1]
L97-L99 · output candidate_representation (return)
C
D
[B, C, D], normalized output

6) Retrieval Head (__call__ + _retrieve_top_k)

L334 · user_representation = build_user_representation(...)
1
D
[B, D] normalized user vector
L365 · scores = matmul(user_representation, corpus_embeddings.T)
N
1
[B, N]
L368 · scores_masked = where(corpus_mask, scores, -INF) (optional)
N
1
[B, N]
L370 · top_k_scores, top_k_indices = jax.lax.top_k(scores, top_k)
K
1
K
1
.
top_k_scores [B, K]
.
top_k_indices [B, K] int
L372 · return (top_k_indices, top_k_scores)
K
1
K
1
.
[B, K]
.
[B, K]
L340-L344 · return RetrievalOutput(...)
1
D
K
1
K
1
.
user_representation [B, D]
.
top_k_indices [B, K]
.
top_k_scores [B, K]