reverse-engineer-objectives

{
  "description": "Reverse-engineer DPO/IPO losses",
  "prompt": "You are reverse-engineering a sealed reference module `reforacle` (importable via `import reforacle`, source unreadable). Implement Python functions reproducing `reforacle`'s EXACT numerical behavior. ONLY dependency may be `torch`. Recover behavior ONLY by calling reforacle and comparing , never decompile.\n\nRun `python3 /app/probe.py` for calling pattern. Oracle called like `reforacle.dpo_loss(pc, pr, rc, rr, beta, label_smoothing)`.\n\nImplement EXACTLY these functions (identical signatures):\n\ndef dpo_loss(pc, pr, rc, rr, beta, label_smoothing):\ndef ipo_loss(pc, pr, rc, rr, beta):\n\nHere pc=policy chosen logprob, pr=policy rejected logprob, rc=reference chosen logprob, rr=reference rejected logprob (each a tensor, per-example, could be (N,) or scalar). Probe extensively:\n\n- dpo_loss: logits = beta * ((pc - rc) - (pr - rr)). Standard DPO loss = -logsigmoid(logits). With label_smoothing: loss = -(1-label_smoothing)*logsigmoid(logits) - label_smoothing*logsigmoid(-logits). DETERMINE exact formula by probing (esp. label_smoothing handling and sign). Determine REDUCTION: does it return per-example (shape N) or mean scalar? Probe with N-element inputs. Determine beta placement.\n- ipo_loss: IPO loss = ((pc-rc) - (pr-rr) - 1/(2*beta))^2. DETERMINE exact: is it (h - 1/(2beta))^2 where h=(pc-rc)-(pr-rr)? Reduction (mean or per-example)? Probe. Some variants: (beta*h - ... ). Nail by crafting inputs with known h and beta.\n\nMethod to determine reduction: pass a vector of N=4 distinct examples; if output is scalar it's reduced (mean), if shape (4,) it's per-example. Then determine if mean or sum by comparing to hand computation.\n\nMethod to determine label_smoothing: set label_smoothing=0 (should reduce to -logsigmoid(logits)) then vary it.\n\nVERIFICATION: Write a test comparing impl to reforacle over >=30 trials each, diverse shapes (scalar, (N,)), diverse beta and label_smoothing values including 0. torch.allclose(atol=1e-6,rtol=1e-5,equal_nan=True). Iterate until all pass.\n\nReturn ONLY final Python source for these 2 functions (no prose/markdown), standalone torch-only (`import torch` assumed). Helpers nested or `_dpo_` prefixed. No reforacle reference in returned code. In your final message alongside code, state the EXACT dpo logits formula, label_smoothing formula, ipo formula, and the reduction convention (per-example vs mean) , I need these for composite objectives.",
  "subagent_type": "general-purpose"
}

, output, 
{"type": "text", "text": "All 60 trials pass across scalar and vector shapes, diverse beta, and label_smoothing (including 0).\n\n```python\ndef dpo_loss(pc, pr, rc, rr, beta, label_smoothing):\n    import torch\n    logits = beta * ((pc - rc) - (pr - rr))\n    loss = (\n        -(1 - label_smoothing) * torch.nn.functional.logsigmoid(logits)\n        - label_smoothing * torch.nn.functional.logsigmoid(-logits)\n    )\n    return loss.mean()\n\n\ndef ipo_loss(pc, pr, rc, rr, beta):\n    h = (pc - rc) - (pr - rr)\n    loss = (h - 1 / (2 * beta)) ** 2\n    return loss.mean()\n```\n\nKey facts for composite objectives:\n\n- DPO logits formula: `logits = beta * ((pc - rc) - (pr - rr))`\n- DPO label_smoothing formula (per-example, before reduction): `-(1 - label_smoothing) * logsigmoid(logits) - label_smoothing * logsigmoid(-logits)`. At `label_smoothing=0` this collapses to `-logsigmoid(logits)`.\n- IPO formula (per-example): `((pc - rc) - (pr - rr) - 1/(2*beta))**2`, i.e. `(h - 1/(2*beta))**2` with `h = (pc-rc)-(pr-rr)`. Beta appears only inside the `1/(2*beta)` margin term, not as a multiplier on `h`.\n- Reduction convention: BOTH functions return a MEAN scalar (shape `[]`), not per-example. Verified by passing N=4 vectors and getting a scalar equal to the per-example `.mean()`.\n\nVerified over 60 trials (scalar and `(N,)` shapes, beta in ~[0.05, 2.05], label_smoothing including 0) with `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)` , all pass."}

{"type": "text", "text": "agentId: ab14fc7203cb72eaf (use SendMessage with to: 'ab14fc7203cb72eaf', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 12305\ntool_uses: 3\nduration_ms: 37779</usage>"}

{"type": "text", "text": "All 60 trials plus edge cases pass with zero failures. Here is the final implementation.\n\n```python\nimport torch\n\n\ndef _lp_masked_mean(values, mask):\n    denom = mask.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return (values * mask).sum() / denom\n\n\ndef logsumexp(x, axis):\n    return torch.logsumexp(x, dim=axis)\n\n\ndef log_softmax(x, axis):\n    return x - torch.logsumexp(x, dim=axis, keepdim=True)\n\n\ndef token_logprobs(logits, labels):\n    lsm = torch.log_softmax(logits, dim=-1)\n    return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n\n\ndef selective_logprobs(logits, labels, mask):\n    lp = token_logprobs(logits, labels)\n    return (lp * mask).sum(dim=-1)\n\n\ndef sequence_logprob(logits, labels, mask, length_normalize):\n    lp = token_logprobs(logits, labels)\n    total = (lp * mask).sum(dim=-1)\n    if length_normalize:\n        denom = mask.sum(dim=-1)\n        safe = torch.where(denom > 0, denom, torch.ones_like(denom))\n        out = total / safe\n        return torch.where(denom > 0, out, torch.zeros_like(out))\n    return total\n\n\ndef logprob_at_temperature(logits, labels, temperature):\n    return token_logprobs(logits / temperature, labels)\n\n\ndef cross_entropy(logits, labels, ignore_index):\n    labels = labels.long()\n    flat_logits = logits.reshape(-1, logits.shape[-1])\n    flat_labels = labels.reshape(-1)\n    lsm = torch.log_softmax(flat_logits, dim=-1)\n    valid = (flat_labels != ignore_index)\n    safe_labels = torch.where(valid, flat_labels, torch.zeros_like(flat_labels))\n    nll = -lsm.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)\n    valid_f = valid.to(nll.dtype)\n    denom = valid_f.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=nll.dtype, device=nll.device)\n    return (nll * valid_f).sum() / denom\n\n\ndef smoothed_nll(logits, labels, smoothing):\n    lsm = torch.log_softmax(logits, dim=-1)\n    nll = -lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n    smooth = -lsm.mean(dim=-1)\n    loss = (1.0 - smoothing) * nll + smoothing * smooth\n    return loss.mean()\n\n\ndef entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    return _lp_masked_mean(ent, mask)\n\n\ndef normalized_entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    vocab = logits.shape[-1]\n    norm = ent / torch.log(torch.tensor(float(vocab), dtype=ent.dtype, device=ent.device))\n    return _lp_masked_mean(norm, mask)\n```\n\nRecovered conventions (all verified over 60 diverse trials + edge cases, 0 failures at atol=1e-6/rtol=1e-5):\n- **logsumexp**: standard, no keepdim; **log_softmax**: `x - logsumexp(x, keepdim=True)`.\n- **token_logprobs**: log_softmax over last dim, gather labels; output shape = labels shape.\n- **selective_logprobs**: `sum(token_logprobs * mask, dim=-1)` , sums over the sequence (last) axis, dropping it.\n- **sequence_logprob**: same masked sum; `length_normalize=True` divides by `mask.sum(-1)`, returning 0 when the mask is all-zero (guarded div).\n- **logprob_at_temperature**: divides logits by temperature before log_softmax.\n- **cross_entropy**: flattens to (-1, V), mean of NLL over non-ignored positions (`sum/count`); returns 0 (not NaN) when all positions are ignored. Works for any ignore_index value.\n- **smoothed_nll**: `mean[(1-s)*nll + s*(-mean_vocab log_softmax)]`.\n- **entropy**: per-position natural-log entropy `-sum(p*log p)`, masked-mean over positions.\n- **normalized_entropy**: entropy / log(vocab_size), masked-mean; returns 0 on empty mask.\n\nNote: `masked_mean` on an all-zero mask returns 0, which is the guard reused by entropy/normalized_entropy."}

{"type": "text", "text": "agentId: acac49472241baa47 (use SendMessage with to: 'acac49472241baa47', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 26571\ntool_uses: 19\nduration_ms: 171626</usage>"}

{"type": "text", "text": "All 1762 tests pass across diverse shapes, group sizes, gamma/lam values, scale_by_std both ways, clip values, and boundary cases.\n\n```python\nimport torch\n\n\ndef _rl_masked_mean(x, mask):\n    return (x * mask).sum() / mask.sum()\n\n\ndef grpo_advantages(rewards, group_size, scale_by_std):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    if scale_by_std:\n        std = groups.std(dim=1, unbiased=False, keepdim=True)\n        adv = adv / (std + 1e-6)\n    return adv.reshape(-1)\n\n\ndef rloo_advantages(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    group_sum = groups.sum(dim=1, keepdim=True)\n    baseline = (group_sum - groups) / (group_size - 1)\n    adv = groups - baseline\n    return adv.reshape(-1)\n\n\ndef group_mean_baseline(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    return adv.reshape(-1)\n\n\ndef gae(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    adv = torch.zeros_like(rewards)\n    gae_acc = torch.zeros_like(rewards[0])\n    next_v = next_value\n    for t in range(T - 1, -1, -1):\n        delta = rewards[t] + gamma * next_v - values[t]\n        gae_acc = delta + gamma * lam * gae_acc\n        adv[t] = gae_acc\n        next_v = values[t]\n    return adv\n\n\ndef discounted_returns(rewards, gamma):\n    T = rewards.shape[0]\n    out = torch.zeros_like(rewards)\n    acc = torch.zeros_like(rewards[0])\n    for t in range(T - 1, -1, -1):\n        acc = rewards[t] + gamma * acc\n        out[t] = acc\n    return out\n\n\ndef lambda_returns(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    out = torch.zeros_like(rewards)\n    next_ret = next_value\n    for t in range(T - 1, -1, -1):\n        next_v = next_value if t == T - 1 else values[t + 1]\n        next_ret = rewards[t] + gamma * ((1 - lam) * next_v + lam * next_ret)\n        out[t] = next_ret\n    return out\n\n\ndef clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):\n    ratio = torch.exp(logp - old_logp)\n    unclipped = ratio * advantages\n    clipped = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) * advantages\n    per = -torch.min(unclipped, clipped)\n    return _rl_masked_mean(per, mask)\n\n\ndef value_loss(values, old_values, returns, clip):\n    unclipped = (values - returns) ** 2\n    vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)\n    clipped = (vpred_clipped - returns) ** 2\n    return 0.5 * torch.max(unclipped, clipped).mean()\n\n\ndef huber_value_loss(values, returns, delta):\n    x = values - returns\n    absx = x.abs()\n    huber = torch.where(absx <= delta, 0.5 * x ** 2, delta * (absx - 0.5 * delta))\n    return huber.mean()\n```\n\nKey conventions you asked about for composite objectives:\n\n- **GAE / returns time-axis & shape**: Time is the **first axis (dim 0)**. Inputs may be 1-D `(T,)` or N-D `(T, ...)` (e.g. `(T, B)`); the recursion runs over dim 0 and per-timestep slices broadcast over trailing dims. `next_value` may be a scalar or shape `(...)` matching the trailing dims. `gae` returns **advantages only** (a single tensor, not a tuple), same shape as `rewards`. `lambda_returns` returns TD(λ) returns and exactly equals `gae + values`. `discounted_returns` also reduces along dim 0.\n\n- **grpo_advantages / group_mean_baseline shape**: Input is flattened then reshaped to `(-1, group_size)` (row-major grouping), and **output is always flat `(N,)`** regardless of input shape. std is **biased** (`unbiased=False`), eps = **1e-6**. Note `group_mean_baseline` returns the **advantages** (`reward - group_mean`), not the raw baseline.\n\n- **clipped_pg_loss sign & clamp**: `ratio = exp(logp - old_logp)`, `clamp(ratio, 1 - clip_low, 1 + clip_high)` (clip_low is the lower bound, clip_high the upper), then **`-min(unclipped, clipped)`** (negative , it is a loss), reduced by **masked mean = `sum(per*mask)/sum(mask)`**.\n\n- **value_loss 0.5 factor**: Yes , `0.5 * mean(max(unclipped, clipped))`, using **max** of clipped vs unclipped squared errors, plain `.mean()` over all elements. `clip` is always numeric (the oracle does not special-case `None`; passing `None` errors). `clip=0.0` degenerates to `0.5*mean((old_values - returns)^2)` vs unclipped max.\n\n- **huber_value_loss**: standard smooth-L1 (`0.5*x^2` for `|x|<=delta`, else `delta*(|x|-0.5*delta)`), reduced by plain mean; equals `F.smooth_l1_loss(..., beta=delta)`."}

{"type": "text", "text": "agentId: ae12d609b60dce0b3 (use SendMessage with to: 'ae12d609b60dce0b3', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 25299\ntool_uses: 11\nduration_ms: 165387</usage>"}

{"type": "text", "text": "All edge cases pass. Implementation complete and verified in `/app/objectives.py`.\n\nFinal Python source for the 6 functions:\n\n```python\nimport torch\n\n\ndef kl_penalty(logp, ref_logp, estimator):\n    logratio = logp - ref_logp\n    if estimator == 'k1':\n        return logratio\n    elif estimator == 'k2':\n        return 0.5 * logratio * logratio\n    elif estimator == 'k3':\n        return torch.exp(-logratio) + logratio - 1\n    else:\n        raise ValueError(estimator)\n\n\ndef reverse_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    return torch.exp(-logratio) + logratio - 1\n\n\ndef symmetric_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    reverse = torch.exp(-logratio) + logratio - 1\n    forward = torch.exp(logratio) - logratio - 1\n    return 0.5 * (reverse + forward)\n\n\ndef importance_ratio(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    if clip is not None:\n        ratio = torch.clamp(ratio, 1 - clip, 1 + clip)\n    return ratio\n\n\ndef clip_fraction(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    return ((ratio > 1 + clip) | (ratio < 1 - clip)).float().mean()\n\n\ndef bradley_terry_logit(chosen_reward, rejected_reward, beta):\n    return beta * (chosen_reward - rejected_reward)\n```\n\nNote on conventions (for composite objectives):\n\n- Supported estimator strings for `kl_penalty`: exactly `'k1'`, `'k2'`, `'k3'`. Any other string raises `ValueError(estimator)`.\n- Sign convention: `logratio = logp - ref_logp` (policy minus reference). Then:\n  - `k1 = logratio` (elementwise, no reduction, same shape as input)\n  - `k2 = 0.5 * logratio**2`\n  - `k3 = exp(-logratio) + logratio - 1` (the non-negative unbiased estimator, using `-logratio` inside exp)\n- `reverse_kl` equals the k3 form: `exp(-logratio) + logratio - 1`. Elementwise, no reduction.\n- `symmetric_kl = 0.5 * (reverse + forward)` where `reverse = exp(-logratio)+logratio-1` and `forward = exp(logratio)-logratio-1`. Elementwise.\n- `importance_ratio`: `ratio = exp(logp - old_logp)`; `clip=None` means no clamp; otherwise clamp to `[1-clip, 1+clip]`. Elementwise.\n- `clip_fraction`: scalar (0-dim tensor) mean of the strict out-of-range condition `(ratio > 1+clip) | (ratio < 1-clip)`. Important: it is NOT `|ratio-1| > clip` , that form disagrees on floating-point boundary values (e.g. `exp(log(1.2))` = 1.20000004 counts as clipped under abs-form but not under the range-form the oracle uses).\n- `bradley_terry_logit = beta * (chosen_reward - rejected_reward)` , a plain logit, no sigmoid/logsigmoid. Elementwise.\n\nVerification: 452 comparisons across 30+ random trials with diverse shapes `(4,), (3,5), (2,3,4), (1,), (10,), (7,2)`, all three estimators, `clip=None` and clip values `{0.0,0.05,0.1,0.2,0.3,0.5,1.0}`, boundary ratios exactly at `1±clip`, and inf/nan inputs , all pass `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)`."}

{"type": "text", "text": "agentId: a5acd56804a9a0d15 (use SendMessage with to: 'a5acd56804a9a0d15', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 29330\ntool_uses: 33\nduration_ms: 212902</usage>"}

{"type": "text", "text": "All 10818 randomized trials pass across diverse shapes (1D/2D/3D), mask kinds (full/zero/sparse/half), all axes, both shift modes, multiple eps values, and empty masks. Combined with the earlier suites and dtype tests, the implementation exactly reproduces `reforacle`.\n\nRecovered conventions:\n- **masked_mean**: `sum(values*mask)/sum(mask)` (mask = float weights), safe-divides to 0 on empty mask; respects `axis` with the reduced dim dropped.\n- **masked_sum**: `sum(values*mask)` over axis (or global).\n- **masked_max**: global scalar, excludes masked positions via `where(mask, values, -inf)`; empty mask returns 0 in the input dtype.\n- **whiten**: masked biased mean/var (÷sum(mask)), `(x-mean)/sqrt(var+1e-8)` applied to all elements; if `shift_mean` is False, adds mean back. No output masking.\n- **masked_whiten**: same as whiten but multiplies the whitened result by mask before optionally adding mean back.\n- **advantage_mean_std**: returns `stack([masked_mean, sqrt(masked_biased_var)])` , biased std, no eps, empty→[0,0].\n- **normalize**: global `(x - x.mean()) / (x.std(unbiased=True) + eps)` , unbiased std, eps added to std (not inside sqrt).\n\nFinal implementation:\n\n```python\nimport torch\n\ndef masked_mean(values, mask, axis=None):\n    if axis is None:\n        s = (values * mask).sum()\n        d = mask.sum()\n    else:\n        s = (values * mask).sum(dim=axis)\n        d = mask.sum(dim=axis)\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef masked_sum(values, mask, axis=None):\n    if axis is None:\n        return (values * mask).sum()\n    return (values * mask).sum(dim=axis)\n\ndef masked_max(values, mask):\n    neg_inf = torch.finfo(values.dtype).min\n    masked_vals = torch.where(mask.bool(), values, torch.full_like(values, neg_inf))\n    if mask.sum() == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return masked_vals.max()\n\ndef _maskprim_safe_mean(values, mask):\n    s = (values * mask).sum()\n    d = mask.sum()\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef _maskprim_mvar(values, mask):\n    mean = _maskprim_safe_mean(values, mask)\n    centered = values - mean\n    var = _maskprim_safe_mean(centered * centered, mask)\n    return mean, var\n\ndef whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef masked_whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    whitened = whitened * mask\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef advantage_mean_std(advantages, mask):\n    mean, var = _maskprim_mvar(advantages, mask)\n    std = torch.sqrt(var)\n    return torch.stack([mean, std])\n\ndef normalize(x, eps):\n    mean = x.mean()\n    std = x.std(unbiased=True)\n    return (x - mean) / (std + eps)\n```"}

{"type": "text", "text": "agentId: af76df31e94c64aa5 (use SendMessage with to: 'af76df31e94c64aa5', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 33906\ntool_uses: 19\nduration_ms: 267714</usage>"}

{"type": "text", "text": "All final edge cases pass, including float64 dtype preservation. The implementation is verified. Here is the final standalone source.\n\n```python\nimport torch\n\n\ndef top_p_mask(probs, p):\n    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)\n    cumsum = torch.cumsum(sorted_probs, dim=-1)\n    cumsum_before = cumsum - sorted_probs\n    keep = cumsum_before < p\n    keep[..., 0] = True\n    kept = sorted_probs * keep.to(sorted_probs.dtype)\n    kept = kept / kept.sum(dim=-1, keepdim=True)\n    out = torch.zeros_like(kept)\n    out.scatter_(-1, sorted_idx, kept)\n    return out\n\n\ndef top_k_mask(logits, k):\n    kth = torch.topk(logits, k, dim=-1).values[..., -1:]\n    return logits >= kth\n\n\ndef argmax_tokens(logits):\n    flipped = torch.flip(logits, dims=[-1])\n    n = logits.shape[-1]\n    idx = torch.argmax(flipped, dim=-1)\n    return (n - 1) - idx\n\n\ndef mode_label(labels):\n    def _mode_1d(v):\n        uniq, counts = torch.unique(v, return_counts=True)\n        max_count = counts.max()\n        candidates = uniq[counts == max_count]\n        return candidates.max()\n\n    if labels.dim() == 1:\n        return _mode_1d(labels)\n    return torch.stack([_mode_1d(labels[i]) for i in range(labels.shape[0])])\n\n\ndef _positional_quantile(x, q, interpolation):\n    xs, _ = torch.sort(x, dim=-1)\n    n = xs.shape[0]\n    pos = (n - 1) * q\n    if interpolation == \"lower\":\n        return xs[int(pos // 1)]\n    lo = int(pos // 1)\n    hi = min(lo + 1, n - 1)\n    frac = pos - lo\n    return xs[lo] * (1 - frac) + xs[hi] * frac\n\n\ndef median_reward(rewards):\n    return _positional_quantile(rewards, 0.5, \"linear\")\n\n\ndef quantile_lower(x, q):\n    return _positional_quantile(x, q, \"lower\")\n\n\ndef pad_mask_from_lengths(lengths, max_len):\n    arange = torch.arange(max_len, device=lengths.device).unsqueeze(0)\n    return (arange < lengths.unsqueeze(1)).to(torch.float32)\n\n\ndef first_nonzero_index(mask):\n    nz = mask != 0\n    any_nz = nz.any(dim=-1)\n    idx = torch.argmax(nz.to(torch.long), dim=-1)\n    return torch.where(any_nz, idx, torch.full_like(idx, -1))\n\n\ndef cumulative_max(x):\n    out = torch.empty_like(x)\n    out[0] = x[0]\n    if x.shape[0] > 1:\n        out[1:] = torch.cummax(x, dim=0)[0][:-1]\n    return out\n\n\ndef bucketize_reward(rewards, boundaries):\n    return torch.bucketize(rewards, boundaries, right=True)\n```\n\n## Summary\n\nI reverse-engineered all 10 functions by probing `reforacle` and verified against 100+ diverse trials (including heavy tie/edge cases) with zero mismatches. Key recovered conventions, several of which are non-obvious:\n\n- **top_p_mask**: nucleus over last dim; keep where `cumsum_before < p` (HF shift, crossing token included), always keeps ≥1 token, output is **renormalized probabilities** in original order (not a boolean mask), zeros elsewhere.\n- **top_k_mask**: **threshold-based** (`logit >= kth-largest value`), so ties at the threshold are all kept and the count may exceed k. Returns bool, last dim.\n- **argmax_tokens**: last dim, but returns the **LAST** maximum index on ties (not torch.argmax's first).\n- **mode_label**: most frequent value per row (dim 0 for 1D, per-row for 2D); tie-break returns the **largest** value (not torch.mode which returns smallest). Handles negatives.\n- **median_reward / quantile_lower**: reduce along dim 0, but with a non-standard rule , **sort along the last dim, then take the positional quantile across dim 0 without re-sorting** (median=linear interp, quantile_lower=floor/lower). Standard `torch.quantile(dim=0)` does NOT reproduce this because it re-sorts each column. For 1D this coincides with normal quantile.\n- **pad_mask_from_lengths**: `(N, max_len)` float32, `arange < length`.\n- **first_nonzero_index**: per-row along last dim, returns **-1** when a row is all zero.\n- **cumulative_max**: exclusive right-shifted running max along **dim 0** (`out[0]=x[0]`, `out[i]=max(x[0..i-1])`).\n- **bucketize_reward**: `torch.bucketize(..., right=True)`.\n\nFiles: `/app/myimpl.py` (final implementation), verification scripts `/app/verify2.py`, `/app/verify4.py`, `/app/verify5.py`. The `/app/objectives.py` scaffold was left unmodified (it contains many other stubs unrelated to this task)."}

{"type": "text", "text": "agentId: aca416a9a99d1536f (use SendMessage with to: 'aca416a9a99d1536f', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 63559\ntool_uses: 37\nduration_ms: 552401</usage>"}

{
  "description": "Reverse-engineer logprob primitives",
  "prompt": "You are reverse-engineering a sealed reference module `reforacle` (importable via `import reforacle`, source unreadable). Implement Python functions reproducing `reforacle`'s EXACT numerical behavior. ONLY dependency may be `torch`. Recover behavior ONLY by calling reforacle and comparing , never decompile.\n\nRun `python3 /app/probe.py` to see calling pattern. Oracle called like `reforacle.log_softmax(x, axis)`.\n\nImplement EXACTLY these functions (identical signatures):\n\ndef logsumexp(x, axis):\ndef log_softmax(x, axis):\ndef selective_logprobs(logits, labels, mask):\ndef token_logprobs(logits, labels):\ndef sequence_logprob(logits, labels, mask, length_normalize):\ndef logprob_at_temperature(logits, labels, temperature):\ndef cross_entropy(logits, labels, ignore_index):\ndef smoothed_nll(logits, labels, smoothing):\ndef entropy(logits, mask):\ndef normalized_entropy(logits, mask):\n\nProbe extensively. Determine EXACT conventions (graded on held-out inputs):\n- logsumexp(x, axis): standard, but confirm axis handling (keepdim? probably not). log_softmax(x,axis): standard x - logsumexp.\n- token_logprobs(logits, labels): logits shape like (..., vocab), labels shape (...). Returns log_softmax(logits).gather(labels). Confirm output shape (labels shape). Determine along which axis softmax is taken (last dim -1). labels dtype long.\n- selective_logprobs(logits, labels, mask): token logprobs but masked. Does it return per-token logprobs * mask, or summed, or masked_mean? Probe output shape carefully. Compare to token_logprobs to see relationship (elementwise mask multiply? zeroed where mask=0? reduced?).\n- sequence_logprob(logits, labels, mask, length_normalize): sum of token logprobs over sequence (masked). length_normalize=True divides by sum(mask) (sequence length), False doesn't. Probe both. Determine reduction axis (last axis = sequence). Output shape.\n- logprob_at_temperature(logits, labels, temperature): applies temperature to logits (logits/temperature) before log_softmax, then gathers labels. Confirm whether temperature divides logits before softmax. Output shape = token logprobs.\n- cross_entropy(logits, labels, ignore_index): standard CE. Determine reduction (mean over non-ignored? sum? none?). How ignore_index is handled (positions with label==ignore_index excluded). Compare to torch.nn.functional.cross_entropy with reduction='mean' and ignore_index. logits shape (N, vocab) or (N, C, ...)? Probe shapes.\n- smoothed_nll(logits, labels, smoothing): label-smoothed negative log likelihood. loss = (1-smoothing)*nll + smoothing*(-mean log_softmax over vocab). Determine exact formula & reduction. Probe smoothing=0 to match plain nll.\n- entropy(logits, mask): entropy of softmax(logits) distribution over vocab: -sum(p*log p) per position, then masked-reduced by mask. Determine: natural log. Reduction = masked_mean over positions? or sum? Probe with mask. Determine relationship to normalized_entropy.\n- normalized_entropy(logits, mask): entropy divided by log(vocab_size)? Probe to find the normalization constant (likely log(V) where V=vocab dim). \n\nVERIFICATION: Write a test comparing your impl to reforacle over >=30 random trials each, diverse shapes (2D (N,V), 3D (B,T,V)), diverse masks, temperature values, smoothing values, ignore_index cases. Use torch.allclose(atol=1e-6,rtol=1e-5,equal_nan=True). For integer-output funcs use exact equality. Iterate until all pass including edge cases.\n\nReturn ONLY final Python source for these 10 functions (no prose/markdown), standalone, torch-only (`import torch` assumed). Helpers as nested functions or `_lp_` prefixed. No reforacle reference in returned code.",
  "subagent_type": "general-purpose"
}

, output, 
{"type": "text", "text": "All 60 trials pass across scalar and vector shapes, diverse beta, and label_smoothing (including 0).\n\n```python\ndef dpo_loss(pc, pr, rc, rr, beta, label_smoothing):\n    import torch\n    logits = beta * ((pc - rc) - (pr - rr))\n    loss = (\n        -(1 - label_smoothing) * torch.nn.functional.logsigmoid(logits)\n        - label_smoothing * torch.nn.functional.logsigmoid(-logits)\n    )\n    return loss.mean()\n\n\ndef ipo_loss(pc, pr, rc, rr, beta):\n    h = (pc - rc) - (pr - rr)\n    loss = (h - 1 / (2 * beta)) ** 2\n    return loss.mean()\n```\n\nKey facts for composite objectives:\n\n- DPO logits formula: `logits = beta * ((pc - rc) - (pr - rr))`\n- DPO label_smoothing formula (per-example, before reduction): `-(1 - label_smoothing) * logsigmoid(logits) - label_smoothing * logsigmoid(-logits)`. At `label_smoothing=0` this collapses to `-logsigmoid(logits)`.\n- IPO formula (per-example): `((pc - rc) - (pr - rr) - 1/(2*beta))**2`, i.e. `(h - 1/(2*beta))**2` with `h = (pc-rc)-(pr-rr)`. Beta appears only inside the `1/(2*beta)` margin term, not as a multiplier on `h`.\n- Reduction convention: BOTH functions return a MEAN scalar (shape `[]`), not per-example. Verified by passing N=4 vectors and getting a scalar equal to the per-example `.mean()`.\n\nVerified over 60 trials (scalar and `(N,)` shapes, beta in ~[0.05, 2.05], label_smoothing including 0) with `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)` , all pass."}

{"type": "text", "text": "agentId: ab14fc7203cb72eaf (use SendMessage with to: 'ab14fc7203cb72eaf', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 12305\ntool_uses: 3\nduration_ms: 37779</usage>"}

{"type": "text", "text": "All 60 trials plus edge cases pass with zero failures. Here is the final implementation.\n\n```python\nimport torch\n\n\ndef _lp_masked_mean(values, mask):\n    denom = mask.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return (values * mask).sum() / denom\n\n\ndef logsumexp(x, axis):\n    return torch.logsumexp(x, dim=axis)\n\n\ndef log_softmax(x, axis):\n    return x - torch.logsumexp(x, dim=axis, keepdim=True)\n\n\ndef token_logprobs(logits, labels):\n    lsm = torch.log_softmax(logits, dim=-1)\n    return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n\n\ndef selective_logprobs(logits, labels, mask):\n    lp = token_logprobs(logits, labels)\n    return (lp * mask).sum(dim=-1)\n\n\ndef sequence_logprob(logits, labels, mask, length_normalize):\n    lp = token_logprobs(logits, labels)\n    total = (lp * mask).sum(dim=-1)\n    if length_normalize:\n        denom = mask.sum(dim=-1)\n        safe = torch.where(denom > 0, denom, torch.ones_like(denom))\n        out = total / safe\n        return torch.where(denom > 0, out, torch.zeros_like(out))\n    return total\n\n\ndef logprob_at_temperature(logits, labels, temperature):\n    return token_logprobs(logits / temperature, labels)\n\n\ndef cross_entropy(logits, labels, ignore_index):\n    labels = labels.long()\n    flat_logits = logits.reshape(-1, logits.shape[-1])\n    flat_labels = labels.reshape(-1)\n    lsm = torch.log_softmax(flat_logits, dim=-1)\n    valid = (flat_labels != ignore_index)\n    safe_labels = torch.where(valid, flat_labels, torch.zeros_like(flat_labels))\n    nll = -lsm.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)\n    valid_f = valid.to(nll.dtype)\n    denom = valid_f.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=nll.dtype, device=nll.device)\n    return (nll * valid_f).sum() / denom\n\n\ndef smoothed_nll(logits, labels, smoothing):\n    lsm = torch.log_softmax(logits, dim=-1)\n    nll = -lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n    smooth = -lsm.mean(dim=-1)\n    loss = (1.0 - smoothing) * nll + smoothing * smooth\n    return loss.mean()\n\n\ndef entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    return _lp_masked_mean(ent, mask)\n\n\ndef normalized_entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    vocab = logits.shape[-1]\n    norm = ent / torch.log(torch.tensor(float(vocab), dtype=ent.dtype, device=ent.device))\n    return _lp_masked_mean(norm, mask)\n```\n\nRecovered conventions (all verified over 60 diverse trials + edge cases, 0 failures at atol=1e-6/rtol=1e-5):\n- **logsumexp**: standard, no keepdim; **log_softmax**: `x - logsumexp(x, keepdim=True)`.\n- **token_logprobs**: log_softmax over last dim, gather labels; output shape = labels shape.\n- **selective_logprobs**: `sum(token_logprobs * mask, dim=-1)` , sums over the sequence (last) axis, dropping it.\n- **sequence_logprob**: same masked sum; `length_normalize=True` divides by `mask.sum(-1)`, returning 0 when the mask is all-zero (guarded div).\n- **logprob_at_temperature**: divides logits by temperature before log_softmax.\n- **cross_entropy**: flattens to (-1, V), mean of NLL over non-ignored positions (`sum/count`); returns 0 (not NaN) when all positions are ignored. Works for any ignore_index value.\n- **smoothed_nll**: `mean[(1-s)*nll + s*(-mean_vocab log_softmax)]`.\n- **entropy**: per-position natural-log entropy `-sum(p*log p)`, masked-mean over positions.\n- **normalized_entropy**: entropy / log(vocab_size), masked-mean; returns 0 on empty mask.\n\nNote: `masked_mean` on an all-zero mask returns 0, which is the guard reused by entropy/normalized_entropy."}

{"type": "text", "text": "agentId: acac49472241baa47 (use SendMessage with to: 'acac49472241baa47', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 26571\ntool_uses: 19\nduration_ms: 171626</usage>"}

{"type": "text", "text": "All 1762 tests pass across diverse shapes, group sizes, gamma/lam values, scale_by_std both ways, clip values, and boundary cases.\n\n```python\nimport torch\n\n\ndef _rl_masked_mean(x, mask):\n    return (x * mask).sum() / mask.sum()\n\n\ndef grpo_advantages(rewards, group_size, scale_by_std):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    if scale_by_std:\n        std = groups.std(dim=1, unbiased=False, keepdim=True)\n        adv = adv / (std + 1e-6)\n    return adv.reshape(-1)\n\n\ndef rloo_advantages(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    group_sum = groups.sum(dim=1, keepdim=True)\n    baseline = (group_sum - groups) / (group_size - 1)\n    adv = groups - baseline\n    return adv.reshape(-1)\n\n\ndef group_mean_baseline(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    return adv.reshape(-1)\n\n\ndef gae(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    adv = torch.zeros_like(rewards)\n    gae_acc = torch.zeros_like(rewards[0])\n    next_v = next_value\n    for t in range(T - 1, -1, -1):\n        delta = rewards[t] + gamma * next_v - values[t]\n        gae_acc = delta + gamma * lam * gae_acc\n        adv[t] = gae_acc\n        next_v = values[t]\n    return adv\n\n\ndef discounted_returns(rewards, gamma):\n    T = rewards.shape[0]\n    out = torch.zeros_like(rewards)\n    acc = torch.zeros_like(rewards[0])\n    for t in range(T - 1, -1, -1):\n        acc = rewards[t] + gamma * acc\n        out[t] = acc\n    return out\n\n\ndef lambda_returns(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    out = torch.zeros_like(rewards)\n    next_ret = next_value\n    for t in range(T - 1, -1, -1):\n        next_v = next_value if t == T - 1 else values[t + 1]\n        next_ret = rewards[t] + gamma * ((1 - lam) * next_v + lam * next_ret)\n        out[t] = next_ret\n    return out\n\n\ndef clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):\n    ratio = torch.exp(logp - old_logp)\n    unclipped = ratio * advantages\n    clipped = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) * advantages\n    per = -torch.min(unclipped, clipped)\n    return _rl_masked_mean(per, mask)\n\n\ndef value_loss(values, old_values, returns, clip):\n    unclipped = (values - returns) ** 2\n    vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)\n    clipped = (vpred_clipped - returns) ** 2\n    return 0.5 * torch.max(unclipped, clipped).mean()\n\n\ndef huber_value_loss(values, returns, delta):\n    x = values - returns\n    absx = x.abs()\n    huber = torch.where(absx <= delta, 0.5 * x ** 2, delta * (absx - 0.5 * delta))\n    return huber.mean()\n```\n\nKey conventions you asked about for composite objectives:\n\n- **GAE / returns time-axis & shape**: Time is the **first axis (dim 0)**. Inputs may be 1-D `(T,)` or N-D `(T, ...)` (e.g. `(T, B)`); the recursion runs over dim 0 and per-timestep slices broadcast over trailing dims. `next_value` may be a scalar or shape `(...)` matching the trailing dims. `gae` returns **advantages only** (a single tensor, not a tuple), same shape as `rewards`. `lambda_returns` returns TD(λ) returns and exactly equals `gae + values`. `discounted_returns` also reduces along dim 0.\n\n- **grpo_advantages / group_mean_baseline shape**: Input is flattened then reshaped to `(-1, group_size)` (row-major grouping), and **output is always flat `(N,)`** regardless of input shape. std is **biased** (`unbiased=False`), eps = **1e-6**. Note `group_mean_baseline` returns the **advantages** (`reward - group_mean`), not the raw baseline.\n\n- **clipped_pg_loss sign & clamp**: `ratio = exp(logp - old_logp)`, `clamp(ratio, 1 - clip_low, 1 + clip_high)` (clip_low is the lower bound, clip_high the upper), then **`-min(unclipped, clipped)`** (negative , it is a loss), reduced by **masked mean = `sum(per*mask)/sum(mask)`**.\n\n- **value_loss 0.5 factor**: Yes , `0.5 * mean(max(unclipped, clipped))`, using **max** of clipped vs unclipped squared errors, plain `.mean()` over all elements. `clip` is always numeric (the oracle does not special-case `None`; passing `None` errors). `clip=0.0` degenerates to `0.5*mean((old_values - returns)^2)` vs unclipped max.\n\n- **huber_value_loss**: standard smooth-L1 (`0.5*x^2` for `|x|<=delta`, else `delta*(|x|-0.5*delta)`), reduced by plain mean; equals `F.smooth_l1_loss(..., beta=delta)`."}

{"type": "text", "text": "agentId: ae12d609b60dce0b3 (use SendMessage with to: 'ae12d609b60dce0b3', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 25299\ntool_uses: 11\nduration_ms: 165387</usage>"}

{"type": "text", "text": "All edge cases pass. Implementation complete and verified in `/app/objectives.py`.\n\nFinal Python source for the 6 functions:\n\n```python\nimport torch\n\n\ndef kl_penalty(logp, ref_logp, estimator):\n    logratio = logp - ref_logp\n    if estimator == 'k1':\n        return logratio\n    elif estimator == 'k2':\n        return 0.5 * logratio * logratio\n    elif estimator == 'k3':\n        return torch.exp(-logratio) + logratio - 1\n    else:\n        raise ValueError(estimator)\n\n\ndef reverse_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    return torch.exp(-logratio) + logratio - 1\n\n\ndef symmetric_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    reverse = torch.exp(-logratio) + logratio - 1\n    forward = torch.exp(logratio) - logratio - 1\n    return 0.5 * (reverse + forward)\n\n\ndef importance_ratio(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    if clip is not None:\n        ratio = torch.clamp(ratio, 1 - clip, 1 + clip)\n    return ratio\n\n\ndef clip_fraction(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    return ((ratio > 1 + clip) | (ratio < 1 - clip)).float().mean()\n\n\ndef bradley_terry_logit(chosen_reward, rejected_reward, beta):\n    return beta * (chosen_reward - rejected_reward)\n```\n\nNote on conventions (for composite objectives):\n\n- Supported estimator strings for `kl_penalty`: exactly `'k1'`, `'k2'`, `'k3'`. Any other string raises `ValueError(estimator)`.\n- Sign convention: `logratio = logp - ref_logp` (policy minus reference). Then:\n  - `k1 = logratio` (elementwise, no reduction, same shape as input)\n  - `k2 = 0.5 * logratio**2`\n  - `k3 = exp(-logratio) + logratio - 1` (the non-negative unbiased estimator, using `-logratio` inside exp)\n- `reverse_kl` equals the k3 form: `exp(-logratio) + logratio - 1`. Elementwise, no reduction.\n- `symmetric_kl = 0.5 * (reverse + forward)` where `reverse = exp(-logratio)+logratio-1` and `forward = exp(logratio)-logratio-1`. Elementwise.\n- `importance_ratio`: `ratio = exp(logp - old_logp)`; `clip=None` means no clamp; otherwise clamp to `[1-clip, 1+clip]`. Elementwise.\n- `clip_fraction`: scalar (0-dim tensor) mean of the strict out-of-range condition `(ratio > 1+clip) | (ratio < 1-clip)`. Important: it is NOT `|ratio-1| > clip` , that form disagrees on floating-point boundary values (e.g. `exp(log(1.2))` = 1.20000004 counts as clipped under abs-form but not under the range-form the oracle uses).\n- `bradley_terry_logit = beta * (chosen_reward - rejected_reward)` , a plain logit, no sigmoid/logsigmoid. Elementwise.\n\nVerification: 452 comparisons across 30+ random trials with diverse shapes `(4,), (3,5), (2,3,4), (1,), (10,), (7,2)`, all three estimators, `clip=None` and clip values `{0.0,0.05,0.1,0.2,0.3,0.5,1.0}`, boundary ratios exactly at `1±clip`, and inf/nan inputs , all pass `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)`."}

{"type": "text", "text": "agentId: a5acd56804a9a0d15 (use SendMessage with to: 'a5acd56804a9a0d15', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 29330\ntool_uses: 33\nduration_ms: 212902</usage>"}

{"type": "text", "text": "All 10818 randomized trials pass across diverse shapes (1D/2D/3D), mask kinds (full/zero/sparse/half), all axes, both shift modes, multiple eps values, and empty masks. Combined with the earlier suites and dtype tests, the implementation exactly reproduces `reforacle`.\n\nRecovered conventions:\n- **masked_mean**: `sum(values*mask)/sum(mask)` (mask = float weights), safe-divides to 0 on empty mask; respects `axis` with the reduced dim dropped.\n- **masked_sum**: `sum(values*mask)` over axis (or global).\n- **masked_max**: global scalar, excludes masked positions via `where(mask, values, -inf)`; empty mask returns 0 in the input dtype.\n- **whiten**: masked biased mean/var (÷sum(mask)), `(x-mean)/sqrt(var+1e-8)` applied to all elements; if `shift_mean` is False, adds mean back. No output masking.\n- **masked_whiten**: same as whiten but multiplies the whitened result by mask before optionally adding mean back.\n- **advantage_mean_std**: returns `stack([masked_mean, sqrt(masked_biased_var)])` , biased std, no eps, empty→[0,0].\n- **normalize**: global `(x - x.mean()) / (x.std(unbiased=True) + eps)` , unbiased std, eps added to std (not inside sqrt).\n\nFinal implementation:\n\n```python\nimport torch\n\ndef masked_mean(values, mask, axis=None):\n    if axis is None:\n        s = (values * mask).sum()\n        d = mask.sum()\n    else:\n        s = (values * mask).sum(dim=axis)\n        d = mask.sum(dim=axis)\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef masked_sum(values, mask, axis=None):\n    if axis is None:\n        return (values * mask).sum()\n    return (values * mask).sum(dim=axis)\n\ndef masked_max(values, mask):\n    neg_inf = torch.finfo(values.dtype).min\n    masked_vals = torch.where(mask.bool(), values, torch.full_like(values, neg_inf))\n    if mask.sum() == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return masked_vals.max()\n\ndef _maskprim_safe_mean(values, mask):\n    s = (values * mask).sum()\n    d = mask.sum()\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef _maskprim_mvar(values, mask):\n    mean = _maskprim_safe_mean(values, mask)\n    centered = values - mean\n    var = _maskprim_safe_mean(centered * centered, mask)\n    return mean, var\n\ndef whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef masked_whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    whitened = whitened * mask\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef advantage_mean_std(advantages, mask):\n    mean, var = _maskprim_mvar(advantages, mask)\n    std = torch.sqrt(var)\n    return torch.stack([mean, std])\n\ndef normalize(x, eps):\n    mean = x.mean()\n    std = x.std(unbiased=True)\n    return (x - mean) / (std + eps)\n```"}

{"type": "text", "text": "agentId: af76df31e94c64aa5 (use SendMessage with to: 'af76df31e94c64aa5', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 33906\ntool_uses: 19\nduration_ms: 267714</usage>"}

{"type": "text", "text": "All final edge cases pass, including float64 dtype preservation. The implementation is verified. Here is the final standalone source.\n\n```python\nimport torch\n\n\ndef top_p_mask(probs, p):\n    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)\n    cumsum = torch.cumsum(sorted_probs, dim=-1)\n    cumsum_before = cumsum - sorted_probs\n    keep = cumsum_before < p\n    keep[..., 0] = True\n    kept = sorted_probs * keep.to(sorted_probs.dtype)\n    kept = kept / kept.sum(dim=-1, keepdim=True)\n    out = torch.zeros_like(kept)\n    out.scatter_(-1, sorted_idx, kept)\n    return out\n\n\ndef top_k_mask(logits, k):\n    kth = torch.topk(logits, k, dim=-1).values[..., -1:]\n    return logits >= kth\n\n\ndef argmax_tokens(logits):\n    flipped = torch.flip(logits, dims=[-1])\n    n = logits.shape[-1]\n    idx = torch.argmax(flipped, dim=-1)\n    return (n - 1) - idx\n\n\ndef mode_label(labels):\n    def _mode_1d(v):\n        uniq, counts = torch.unique(v, return_counts=True)\n        max_count = counts.max()\n        candidates = uniq[counts == max_count]\n        return candidates.max()\n\n    if labels.dim() == 1:\n        return _mode_1d(labels)\n    return torch.stack([_mode_1d(labels[i]) for i in range(labels.shape[0])])\n\n\ndef _positional_quantile(x, q, interpolation):\n    xs, _ = torch.sort(x, dim=-1)\n    n = xs.shape[0]\n    pos = (n - 1) * q\n    if interpolation == \"lower\":\n        return xs[int(pos // 1)]\n    lo = int(pos // 1)\n    hi = min(lo + 1, n - 1)\n    frac = pos - lo\n    return xs[lo] * (1 - frac) + xs[hi] * frac\n\n\ndef median_reward(rewards):\n    return _positional_quantile(rewards, 0.5, \"linear\")\n\n\ndef quantile_lower(x, q):\n    return _positional_quantile(x, q, \"lower\")\n\n\ndef pad_mask_from_lengths(lengths, max_len):\n    arange = torch.arange(max_len, device=lengths.device).unsqueeze(0)\n    return (arange < lengths.unsqueeze(1)).to(torch.float32)\n\n\ndef first_nonzero_index(mask):\n    nz = mask != 0\n    any_nz = nz.any(dim=-1)\n    idx = torch.argmax(nz.to(torch.long), dim=-1)\n    return torch.where(any_nz, idx, torch.full_like(idx, -1))\n\n\ndef cumulative_max(x):\n    out = torch.empty_like(x)\n    out[0] = x[0]\n    if x.shape[0] > 1:\n        out[1:] = torch.cummax(x, dim=0)[0][:-1]\n    return out\n\n\ndef bucketize_reward(rewards, boundaries):\n    return torch.bucketize(rewards, boundaries, right=True)\n```\n\n## Summary\n\nI reverse-engineered all 10 functions by probing `reforacle` and verified against 100+ diverse trials (including heavy tie/edge cases) with zero mismatches. Key recovered conventions, several of which are non-obvious:\n\n- **top_p_mask**: nucleus over last dim; keep where `cumsum_before < p` (HF shift, crossing token included), always keeps ≥1 token, output is **renormalized probabilities** in original order (not a boolean mask), zeros elsewhere.\n- **top_k_mask**: **threshold-based** (`logit >= kth-largest value`), so ties at the threshold are all kept and the count may exceed k. Returns bool, last dim.\n- **argmax_tokens**: last dim, but returns the **LAST** maximum index on ties (not torch.argmax's first).\n- **mode_label**: most frequent value per row (dim 0 for 1D, per-row for 2D); tie-break returns the **largest** value (not torch.mode which returns smallest). Handles negatives.\n- **median_reward / quantile_lower**: reduce along dim 0, but with a non-standard rule , **sort along the last dim, then take the positional quantile across dim 0 without re-sorting** (median=linear interp, quantile_lower=floor/lower). Standard `torch.quantile(dim=0)` does NOT reproduce this because it re-sorts each column. For 1D this coincides with normal quantile.\n- **pad_mask_from_lengths**: `(N, max_len)` float32, `arange < length`.\n- **first_nonzero_index**: per-row along last dim, returns **-1** when a row is all zero.\n- **cumulative_max**: exclusive right-shifted running max along **dim 0** (`out[0]=x[0]`, `out[i]=max(x[0..i-1])`).\n- **bucketize_reward**: `torch.bucketize(..., right=True)`.\n\nFiles: `/app/myimpl.py` (final implementation), verification scripts `/app/verify2.py`, `/app/verify4.py`, `/app/verify5.py`. The `/app/objectives.py` scaffold was left unmodified (it contains many other stubs unrelated to this task)."}

{"type": "text", "text": "agentId: aca416a9a99d1536f (use SendMessage with to: 'aca416a9a99d1536f', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 63559\ntool_uses: 37\nduration_ms: 552401</usage>"}

{
  "description": "Reverse-engineer RL advantage/value funcs",
  "prompt": "You are reverse-engineering a sealed reference module `reforacle` (importable via `import reforacle`, source unreadable). Implement Python functions reproducing `reforacle`'s EXACT numerical behavior. ONLY dependency may be `torch`. Recover behavior ONLY by calling reforacle and comparing , never decompile.\n\nRun `python3 /app/probe.py` for calling pattern. Oracle called like `reforacle.gae(rewards, values, next_value, gamma, lam)`.\n\nImplement EXACTLY these functions (identical signatures):\n\ndef grpo_advantages(rewards, group_size, scale_by_std):\ndef rloo_advantages(rewards, group_size):\ndef group_mean_baseline(rewards, group_size):\ndef gae(rewards, values, next_value, gamma, lam):\ndef discounted_returns(rewards, gamma):\ndef lambda_returns(rewards, values, next_value, gamma, lam):\ndef clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):\ndef value_loss(values, old_values, returns, clip):\ndef huber_value_loss(values, returns, delta):\n\nProbe extensively for exact conventions (graded on held-out inputs):\n\n- grpo_advantages(rewards, group_size, scale_by_std): rewards is a flat tensor of shape (N,) where N is divisible by group_size (or shape (num_groups, group_size)? PROBE the expected shape). Group rewards into groups; advantage = reward - group_mean; if scale_by_std=True divide by (group_std + eps). Determine: group std unbiased or biased? eps value (probe, likely 1e-4 or 1e-6)? Output shape (same as input)? Probe scale_by_std True and False. Determine how groups are formed (reshape N->(N/group_size, group_size)).\n- rloo_advantages(rewards, group_size): leave-one-out baseline. advantage_i = reward_i - mean(other rewards in group) = reward_i - (group_sum - reward_i)/(group_size-1). Confirm exact formula and shape.\n- group_mean_baseline(rewards, group_size): returns baseline = per-group mean broadcast to each element? Or the advantages (reward-mean)? Probe: likely returns the group-mean baseline (same shape as rewards, each element = its group's mean). Confirm.\n- gae(rewards, values, next_value, gamma, lam): Generalized Advantage Estimation. rewards, values shape (T,) or (B,T)? PROBE. next_value is bootstrap value after last step. delta_t = r_t + gamma*V_{t+1} - V_t (V_T = next_value). A_t = delta_t + gamma*lam*A_{t+1}. Returns advantages (shape of rewards). Determine time axis (last dim). Probe carefully with T=3 crafted values to verify recursion & whether it returns advantages or (advantages, returns).\n- discounted_returns(rewards, gamma): G_t = r_t + gamma*G_{t+1}, backward cumulative. No bootstrap. Shape of rewards. Probe axis.\n- lambda_returns(rewards, values, next_value, gamma, lam): TD(lambda) returns = gae advantages + values, OR direct recursion R_t = r_t + gamma*((1-lam)*V_{t+1} + lam*R_{t+1}). Determine exact and whether equals gae+values. Probe.\n- clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high): PPO clipped policy-gradient loss. ratio=exp(logp-old_logp). unclipped = ratio*adv; clipped = clamp(ratio, 1-clip_low, 1+clip_high)*adv. loss = -mean over masked positions of min(unclipped, clipped). DETERMINE: sign (negative for loss?), the clamp bounds (1-clip_low and 1+clip_high , confirm which is low/high), reduction (masked_mean using mask). Probe clip_low, clip_high possibly different. Confirm min() then negate then masked-mean. Determine masked reduction convention (divide by sum(mask)).\n- value_loss(values, old_values, returns, clip): PPO value loss. unclipped=(values-returns)^2; vpred_clipped = old_values + clamp(values-old_values, -clip, clip); clipped=(vpred_clipped-returns)^2; loss = 0.5*mean(max(unclipped,clipped))? DETERMINE the 0.5 factor, max vs no-max, reduction (mean over all), and clip=None meaning no clipping. Probe.\n- huber_value_loss(values, returns, delta): Huber/smooth-L1 loss between values and returns with threshold delta. Determine reduction (mean) and exact Huber formula (0.5*x^2 if |x|<=delta else delta*(|x|-0.5*delta)). Probe delta boundary.\n\nVERIFICATION: Write a test comparing impl to reforacle over >=30 trials each, diverse shapes, group_size divisors, gamma/lam values, clip None and values, scale_by_std both, boundary cases. torch.allclose(atol=1e-6,rtol=1e-5,equal_nan=True). Iterate until all pass.\n\nReturn ONLY final Python source for these 9 functions (no prose/markdown), standalone torch-only (`import torch` assumed). Helpers nested or `_rl_` prefixed. No reforacle reference in returned code. In your final message alongside code, note the exact GAE/returns time-axis and shape convention, the clipped_pg_loss sign & clamp convention, and value_loss 0.5 factor , I need these for composite objectives.",
  "subagent_type": "general-purpose"
}

, output, 
{"type": "text", "text": "All 60 trials pass across scalar and vector shapes, diverse beta, and label_smoothing (including 0).\n\n```python\ndef dpo_loss(pc, pr, rc, rr, beta, label_smoothing):\n    import torch\n    logits = beta * ((pc - rc) - (pr - rr))\n    loss = (\n        -(1 - label_smoothing) * torch.nn.functional.logsigmoid(logits)\n        - label_smoothing * torch.nn.functional.logsigmoid(-logits)\n    )\n    return loss.mean()\n\n\ndef ipo_loss(pc, pr, rc, rr, beta):\n    h = (pc - rc) - (pr - rr)\n    loss = (h - 1 / (2 * beta)) ** 2\n    return loss.mean()\n```\n\nKey facts for composite objectives:\n\n- DPO logits formula: `logits = beta * ((pc - rc) - (pr - rr))`\n- DPO label_smoothing formula (per-example, before reduction): `-(1 - label_smoothing) * logsigmoid(logits) - label_smoothing * logsigmoid(-logits)`. At `label_smoothing=0` this collapses to `-logsigmoid(logits)`.\n- IPO formula (per-example): `((pc - rc) - (pr - rr) - 1/(2*beta))**2`, i.e. `(h - 1/(2*beta))**2` with `h = (pc-rc)-(pr-rr)`. Beta appears only inside the `1/(2*beta)` margin term, not as a multiplier on `h`.\n- Reduction convention: BOTH functions return a MEAN scalar (shape `[]`), not per-example. Verified by passing N=4 vectors and getting a scalar equal to the per-example `.mean()`.\n\nVerified over 60 trials (scalar and `(N,)` shapes, beta in ~[0.05, 2.05], label_smoothing including 0) with `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)` , all pass."}

{"type": "text", "text": "agentId: ab14fc7203cb72eaf (use SendMessage with to: 'ab14fc7203cb72eaf', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 12305\ntool_uses: 3\nduration_ms: 37779</usage>"}

{"type": "text", "text": "All 60 trials plus edge cases pass with zero failures. Here is the final implementation.\n\n```python\nimport torch\n\n\ndef _lp_masked_mean(values, mask):\n    denom = mask.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return (values * mask).sum() / denom\n\n\ndef logsumexp(x, axis):\n    return torch.logsumexp(x, dim=axis)\n\n\ndef log_softmax(x, axis):\n    return x - torch.logsumexp(x, dim=axis, keepdim=True)\n\n\ndef token_logprobs(logits, labels):\n    lsm = torch.log_softmax(logits, dim=-1)\n    return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n\n\ndef selective_logprobs(logits, labels, mask):\n    lp = token_logprobs(logits, labels)\n    return (lp * mask).sum(dim=-1)\n\n\ndef sequence_logprob(logits, labels, mask, length_normalize):\n    lp = token_logprobs(logits, labels)\n    total = (lp * mask).sum(dim=-1)\n    if length_normalize:\n        denom = mask.sum(dim=-1)\n        safe = torch.where(denom > 0, denom, torch.ones_like(denom))\n        out = total / safe\n        return torch.where(denom > 0, out, torch.zeros_like(out))\n    return total\n\n\ndef logprob_at_temperature(logits, labels, temperature):\n    return token_logprobs(logits / temperature, labels)\n\n\ndef cross_entropy(logits, labels, ignore_index):\n    labels = labels.long()\n    flat_logits = logits.reshape(-1, logits.shape[-1])\n    flat_labels = labels.reshape(-1)\n    lsm = torch.log_softmax(flat_logits, dim=-1)\n    valid = (flat_labels != ignore_index)\n    safe_labels = torch.where(valid, flat_labels, torch.zeros_like(flat_labels))\n    nll = -lsm.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)\n    valid_f = valid.to(nll.dtype)\n    denom = valid_f.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=nll.dtype, device=nll.device)\n    return (nll * valid_f).sum() / denom\n\n\ndef smoothed_nll(logits, labels, smoothing):\n    lsm = torch.log_softmax(logits, dim=-1)\n    nll = -lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n    smooth = -lsm.mean(dim=-1)\n    loss = (1.0 - smoothing) * nll + smoothing * smooth\n    return loss.mean()\n\n\ndef entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    return _lp_masked_mean(ent, mask)\n\n\ndef normalized_entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    vocab = logits.shape[-1]\n    norm = ent / torch.log(torch.tensor(float(vocab), dtype=ent.dtype, device=ent.device))\n    return _lp_masked_mean(norm, mask)\n```\n\nRecovered conventions (all verified over 60 diverse trials + edge cases, 0 failures at atol=1e-6/rtol=1e-5):\n- **logsumexp**: standard, no keepdim; **log_softmax**: `x - logsumexp(x, keepdim=True)`.\n- **token_logprobs**: log_softmax over last dim, gather labels; output shape = labels shape.\n- **selective_logprobs**: `sum(token_logprobs * mask, dim=-1)` , sums over the sequence (last) axis, dropping it.\n- **sequence_logprob**: same masked sum; `length_normalize=True` divides by `mask.sum(-1)`, returning 0 when the mask is all-zero (guarded div).\n- **logprob_at_temperature**: divides logits by temperature before log_softmax.\n- **cross_entropy**: flattens to (-1, V), mean of NLL over non-ignored positions (`sum/count`); returns 0 (not NaN) when all positions are ignored. Works for any ignore_index value.\n- **smoothed_nll**: `mean[(1-s)*nll + s*(-mean_vocab log_softmax)]`.\n- **entropy**: per-position natural-log entropy `-sum(p*log p)`, masked-mean over positions.\n- **normalized_entropy**: entropy / log(vocab_size), masked-mean; returns 0 on empty mask.\n\nNote: `masked_mean` on an all-zero mask returns 0, which is the guard reused by entropy/normalized_entropy."}

{"type": "text", "text": "agentId: acac49472241baa47 (use SendMessage with to: 'acac49472241baa47', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 26571\ntool_uses: 19\nduration_ms: 171626</usage>"}

{"type": "text", "text": "All 1762 tests pass across diverse shapes, group sizes, gamma/lam values, scale_by_std both ways, clip values, and boundary cases.\n\n```python\nimport torch\n\n\ndef _rl_masked_mean(x, mask):\n    return (x * mask).sum() / mask.sum()\n\n\ndef grpo_advantages(rewards, group_size, scale_by_std):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    if scale_by_std:\n        std = groups.std(dim=1, unbiased=False, keepdim=True)\n        adv = adv / (std + 1e-6)\n    return adv.reshape(-1)\n\n\ndef rloo_advantages(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    group_sum = groups.sum(dim=1, keepdim=True)\n    baseline = (group_sum - groups) / (group_size - 1)\n    adv = groups - baseline\n    return adv.reshape(-1)\n\n\ndef group_mean_baseline(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    return adv.reshape(-1)\n\n\ndef gae(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    adv = torch.zeros_like(rewards)\n    gae_acc = torch.zeros_like(rewards[0])\n    next_v = next_value\n    for t in range(T - 1, -1, -1):\n        delta = rewards[t] + gamma * next_v - values[t]\n        gae_acc = delta + gamma * lam * gae_acc\n        adv[t] = gae_acc\n        next_v = values[t]\n    return adv\n\n\ndef discounted_returns(rewards, gamma):\n    T = rewards.shape[0]\n    out = torch.zeros_like(rewards)\n    acc = torch.zeros_like(rewards[0])\n    for t in range(T - 1, -1, -1):\n        acc = rewards[t] + gamma * acc\n        out[t] = acc\n    return out\n\n\ndef lambda_returns(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    out = torch.zeros_like(rewards)\n    next_ret = next_value\n    for t in range(T - 1, -1, -1):\n        next_v = next_value if t == T - 1 else values[t + 1]\n        next_ret = rewards[t] + gamma * ((1 - lam) * next_v + lam * next_ret)\n        out[t] = next_ret\n    return out\n\n\ndef clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):\n    ratio = torch.exp(logp - old_logp)\n    unclipped = ratio * advantages\n    clipped = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) * advantages\n    per = -torch.min(unclipped, clipped)\n    return _rl_masked_mean(per, mask)\n\n\ndef value_loss(values, old_values, returns, clip):\n    unclipped = (values - returns) ** 2\n    vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)\n    clipped = (vpred_clipped - returns) ** 2\n    return 0.5 * torch.max(unclipped, clipped).mean()\n\n\ndef huber_value_loss(values, returns, delta):\n    x = values - returns\n    absx = x.abs()\n    huber = torch.where(absx <= delta, 0.5 * x ** 2, delta * (absx - 0.5 * delta))\n    return huber.mean()\n```\n\nKey conventions you asked about for composite objectives:\n\n- **GAE / returns time-axis & shape**: Time is the **first axis (dim 0)**. Inputs may be 1-D `(T,)` or N-D `(T, ...)` (e.g. `(T, B)`); the recursion runs over dim 0 and per-timestep slices broadcast over trailing dims. `next_value` may be a scalar or shape `(...)` matching the trailing dims. `gae` returns **advantages only** (a single tensor, not a tuple), same shape as `rewards`. `lambda_returns` returns TD(λ) returns and exactly equals `gae + values`. `discounted_returns` also reduces along dim 0.\n\n- **grpo_advantages / group_mean_baseline shape**: Input is flattened then reshaped to `(-1, group_size)` (row-major grouping), and **output is always flat `(N,)`** regardless of input shape. std is **biased** (`unbiased=False`), eps = **1e-6**. Note `group_mean_baseline` returns the **advantages** (`reward - group_mean`), not the raw baseline.\n\n- **clipped_pg_loss sign & clamp**: `ratio = exp(logp - old_logp)`, `clamp(ratio, 1 - clip_low, 1 + clip_high)` (clip_low is the lower bound, clip_high the upper), then **`-min(unclipped, clipped)`** (negative , it is a loss), reduced by **masked mean = `sum(per*mask)/sum(mask)`**.\n\n- **value_loss 0.5 factor**: Yes , `0.5 * mean(max(unclipped, clipped))`, using **max** of clipped vs unclipped squared errors, plain `.mean()` over all elements. `clip` is always numeric (the oracle does not special-case `None`; passing `None` errors). `clip=0.0` degenerates to `0.5*mean((old_values - returns)^2)` vs unclipped max.\n\n- **huber_value_loss**: standard smooth-L1 (`0.5*x^2` for `|x|<=delta`, else `delta*(|x|-0.5*delta)`), reduced by plain mean; equals `F.smooth_l1_loss(..., beta=delta)`."}

{"type": "text", "text": "agentId: ae12d609b60dce0b3 (use SendMessage with to: 'ae12d609b60dce0b3', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 25299\ntool_uses: 11\nduration_ms: 165387</usage>"}

{"type": "text", "text": "All edge cases pass. Implementation complete and verified in `/app/objectives.py`.\n\nFinal Python source for the 6 functions:\n\n```python\nimport torch\n\n\ndef kl_penalty(logp, ref_logp, estimator):\n    logratio = logp - ref_logp\n    if estimator == 'k1':\n        return logratio\n    elif estimator == 'k2':\n        return 0.5 * logratio * logratio\n    elif estimator == 'k3':\n        return torch.exp(-logratio) + logratio - 1\n    else:\n        raise ValueError(estimator)\n\n\ndef reverse_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    return torch.exp(-logratio) + logratio - 1\n\n\ndef symmetric_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    reverse = torch.exp(-logratio) + logratio - 1\n    forward = torch.exp(logratio) - logratio - 1\n    return 0.5 * (reverse + forward)\n\n\ndef importance_ratio(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    if clip is not None:\n        ratio = torch.clamp(ratio, 1 - clip, 1 + clip)\n    return ratio\n\n\ndef clip_fraction(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    return ((ratio > 1 + clip) | (ratio < 1 - clip)).float().mean()\n\n\ndef bradley_terry_logit(chosen_reward, rejected_reward, beta):\n    return beta * (chosen_reward - rejected_reward)\n```\n\nNote on conventions (for composite objectives):\n\n- Supported estimator strings for `kl_penalty`: exactly `'k1'`, `'k2'`, `'k3'`. Any other string raises `ValueError(estimator)`.\n- Sign convention: `logratio = logp - ref_logp` (policy minus reference). Then:\n  - `k1 = logratio` (elementwise, no reduction, same shape as input)\n  - `k2 = 0.5 * logratio**2`\n  - `k3 = exp(-logratio) + logratio - 1` (the non-negative unbiased estimator, using `-logratio` inside exp)\n- `reverse_kl` equals the k3 form: `exp(-logratio) + logratio - 1`. Elementwise, no reduction.\n- `symmetric_kl = 0.5 * (reverse + forward)` where `reverse = exp(-logratio)+logratio-1` and `forward = exp(logratio)-logratio-1`. Elementwise.\n- `importance_ratio`: `ratio = exp(logp - old_logp)`; `clip=None` means no clamp; otherwise clamp to `[1-clip, 1+clip]`. Elementwise.\n- `clip_fraction`: scalar (0-dim tensor) mean of the strict out-of-range condition `(ratio > 1+clip) | (ratio < 1-clip)`. Important: it is NOT `|ratio-1| > clip` , that form disagrees on floating-point boundary values (e.g. `exp(log(1.2))` = 1.20000004 counts as clipped under abs-form but not under the range-form the oracle uses).\n- `bradley_terry_logit = beta * (chosen_reward - rejected_reward)` , a plain logit, no sigmoid/logsigmoid. Elementwise.\n\nVerification: 452 comparisons across 30+ random trials with diverse shapes `(4,), (3,5), (2,3,4), (1,), (10,), (7,2)`, all three estimators, `clip=None` and clip values `{0.0,0.05,0.1,0.2,0.3,0.5,1.0}`, boundary ratios exactly at `1±clip`, and inf/nan inputs , all pass `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)`."}

{"type": "text", "text": "agentId: a5acd56804a9a0d15 (use SendMessage with to: 'a5acd56804a9a0d15', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 29330\ntool_uses: 33\nduration_ms: 212902</usage>"}

{"type": "text", "text": "All 10818 randomized trials pass across diverse shapes (1D/2D/3D), mask kinds (full/zero/sparse/half), all axes, both shift modes, multiple eps values, and empty masks. Combined with the earlier suites and dtype tests, the implementation exactly reproduces `reforacle`.\n\nRecovered conventions:\n- **masked_mean**: `sum(values*mask)/sum(mask)` (mask = float weights), safe-divides to 0 on empty mask; respects `axis` with the reduced dim dropped.\n- **masked_sum**: `sum(values*mask)` over axis (or global).\n- **masked_max**: global scalar, excludes masked positions via `where(mask, values, -inf)`; empty mask returns 0 in the input dtype.\n- **whiten**: masked biased mean/var (÷sum(mask)), `(x-mean)/sqrt(var+1e-8)` applied to all elements; if `shift_mean` is False, adds mean back. No output masking.\n- **masked_whiten**: same as whiten but multiplies the whitened result by mask before optionally adding mean back.\n- **advantage_mean_std**: returns `stack([masked_mean, sqrt(masked_biased_var)])` , biased std, no eps, empty→[0,0].\n- **normalize**: global `(x - x.mean()) / (x.std(unbiased=True) + eps)` , unbiased std, eps added to std (not inside sqrt).\n\nFinal implementation:\n\n```python\nimport torch\n\ndef masked_mean(values, mask, axis=None):\n    if axis is None:\n        s = (values * mask).sum()\n        d = mask.sum()\n    else:\n        s = (values * mask).sum(dim=axis)\n        d = mask.sum(dim=axis)\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef masked_sum(values, mask, axis=None):\n    if axis is None:\n        return (values * mask).sum()\n    return (values * mask).sum(dim=axis)\n\ndef masked_max(values, mask):\n    neg_inf = torch.finfo(values.dtype).min\n    masked_vals = torch.where(mask.bool(), values, torch.full_like(values, neg_inf))\n    if mask.sum() == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return masked_vals.max()\n\ndef _maskprim_safe_mean(values, mask):\n    s = (values * mask).sum()\n    d = mask.sum()\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef _maskprim_mvar(values, mask):\n    mean = _maskprim_safe_mean(values, mask)\n    centered = values - mean\n    var = _maskprim_safe_mean(centered * centered, mask)\n    return mean, var\n\ndef whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef masked_whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    whitened = whitened * mask\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef advantage_mean_std(advantages, mask):\n    mean, var = _maskprim_mvar(advantages, mask)\n    std = torch.sqrt(var)\n    return torch.stack([mean, std])\n\ndef normalize(x, eps):\n    mean = x.mean()\n    std = x.std(unbiased=True)\n    return (x - mean) / (std + eps)\n```"}

{"type": "text", "text": "agentId: af76df31e94c64aa5 (use SendMessage with to: 'af76df31e94c64aa5', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 33906\ntool_uses: 19\nduration_ms: 267714</usage>"}

{"type": "text", "text": "All final edge cases pass, including float64 dtype preservation. The implementation is verified. Here is the final standalone source.\n\n```python\nimport torch\n\n\ndef top_p_mask(probs, p):\n    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)\n    cumsum = torch.cumsum(sorted_probs, dim=-1)\n    cumsum_before = cumsum - sorted_probs\n    keep = cumsum_before < p\n    keep[..., 0] = True\n    kept = sorted_probs * keep.to(sorted_probs.dtype)\n    kept = kept / kept.sum(dim=-1, keepdim=True)\n    out = torch.zeros_like(kept)\n    out.scatter_(-1, sorted_idx, kept)\n    return out\n\n\ndef top_k_mask(logits, k):\n    kth = torch.topk(logits, k, dim=-1).values[..., -1:]\n    return logits >= kth\n\n\ndef argmax_tokens(logits):\n    flipped = torch.flip(logits, dims=[-1])\n    n = logits.shape[-1]\n    idx = torch.argmax(flipped, dim=-1)\n    return (n - 1) - idx\n\n\ndef mode_label(labels):\n    def _mode_1d(v):\n        uniq, counts = torch.unique(v, return_counts=True)\n        max_count = counts.max()\n        candidates = uniq[counts == max_count]\n        return candidates.max()\n\n    if labels.dim() == 1:\n        return _mode_1d(labels)\n    return torch.stack([_mode_1d(labels[i]) for i in range(labels.shape[0])])\n\n\ndef _positional_quantile(x, q, interpolation):\n    xs, _ = torch.sort(x, dim=-1)\n    n = xs.shape[0]\n    pos = (n - 1) * q\n    if interpolation == \"lower\":\n        return xs[int(pos // 1)]\n    lo = int(pos // 1)\n    hi = min(lo + 1, n - 1)\n    frac = pos - lo\n    return xs[lo] * (1 - frac) + xs[hi] * frac\n\n\ndef median_reward(rewards):\n    return _positional_quantile(rewards, 0.5, \"linear\")\n\n\ndef quantile_lower(x, q):\n    return _positional_quantile(x, q, \"lower\")\n\n\ndef pad_mask_from_lengths(lengths, max_len):\n    arange = torch.arange(max_len, device=lengths.device).unsqueeze(0)\n    return (arange < lengths.unsqueeze(1)).to(torch.float32)\n\n\ndef first_nonzero_index(mask):\n    nz = mask != 0\n    any_nz = nz.any(dim=-1)\n    idx = torch.argmax(nz.to(torch.long), dim=-1)\n    return torch.where(any_nz, idx, torch.full_like(idx, -1))\n\n\ndef cumulative_max(x):\n    out = torch.empty_like(x)\n    out[0] = x[0]\n    if x.shape[0] > 1:\n        out[1:] = torch.cummax(x, dim=0)[0][:-1]\n    return out\n\n\ndef bucketize_reward(rewards, boundaries):\n    return torch.bucketize(rewards, boundaries, right=True)\n```\n\n## Summary\n\nI reverse-engineered all 10 functions by probing `reforacle` and verified against 100+ diverse trials (including heavy tie/edge cases) with zero mismatches. Key recovered conventions, several of which are non-obvious:\n\n- **top_p_mask**: nucleus over last dim; keep where `cumsum_before < p` (HF shift, crossing token included), always keeps ≥1 token, output is **renormalized probabilities** in original order (not a boolean mask), zeros elsewhere.\n- **top_k_mask**: **threshold-based** (`logit >= kth-largest value`), so ties at the threshold are all kept and the count may exceed k. Returns bool, last dim.\n- **argmax_tokens**: last dim, but returns the **LAST** maximum index on ties (not torch.argmax's first).\n- **mode_label**: most frequent value per row (dim 0 for 1D, per-row for 2D); tie-break returns the **largest** value (not torch.mode which returns smallest). Handles negatives.\n- **median_reward / quantile_lower**: reduce along dim 0, but with a non-standard rule , **sort along the last dim, then take the positional quantile across dim 0 without re-sorting** (median=linear interp, quantile_lower=floor/lower). Standard `torch.quantile(dim=0)` does NOT reproduce this because it re-sorts each column. For 1D this coincides with normal quantile.\n- **pad_mask_from_lengths**: `(N, max_len)` float32, `arange < length`.\n- **first_nonzero_index**: per-row along last dim, returns **-1** when a row is all zero.\n- **cumulative_max**: exclusive right-shifted running max along **dim 0** (`out[0]=x[0]`, `out[i]=max(x[0..i-1])`).\n- **bucketize_reward**: `torch.bucketize(..., right=True)`.\n\nFiles: `/app/myimpl.py` (final implementation), verification scripts `/app/verify2.py`, `/app/verify4.py`, `/app/verify5.py`. The `/app/objectives.py` scaffold was left unmodified (it contains many other stubs unrelated to this task)."}

{"type": "text", "text": "agentId: aca416a9a99d1536f (use SendMessage with to: 'aca416a9a99d1536f', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 63559\ntool_uses: 37\nduration_ms: 552401</usage>"}

{
  "description": "Reverse-engineer KL and ratio funcs",
  "prompt": "You are reverse-engineering a sealed reference module `reforacle` (importable via `import reforacle`, source unreadable). Implement Python functions reproducing `reforacle`'s EXACT numerical behavior. ONLY dependency may be `torch`. Recover behavior ONLY by calling reforacle and comparing , never decompile.\n\nRun `python3 /app/probe.py` for calling pattern. Oracle called like `reforacle.kl_penalty(logp, ref_logp, estimator)`.\n\nImplement EXACTLY these functions (identical signatures):\n\ndef kl_penalty(logp, ref_logp, estimator):\ndef reverse_kl(logp, ref_logp):\ndef symmetric_kl(logp, ref_logp):\ndef importance_ratio(logp, old_logp, clip):\ndef clip_fraction(logp, old_logp, clip):\ndef bradley_terry_logit(chosen_reward, rejected_reward, beta):\n\nHere logp/ref_logp/old_logp are LOG-PROBABILITIES (per-token or per-sequence tensors), elementwise. Probe extensively for exact conventions:\n\n- kl_penalty(logp, ref_logp, estimator): estimator is a STRING selecting the KL estimator. Determine which string values are supported by probing: try 'k1', 'k2', 'k3', 'kl', 'abs', 'mse', 'j' etc. (call and see which don't error). Standard (John Schulman) estimators between policy logp and ref ref_logp, let r = ref_logp - logp (log ratio) OR logp - ref_logp , DETERMINE sign by probing:\n  * k1 = -logratio  or  (logp - ref_logp)? Probe.\n  * k2 = 0.5 * (logratio)^2\n  * k3 = exp(logratio) - logratio - 1  (>=0 unbiased estimator)\n  Determine the EXACT definition and sign of logratio for each supported estimator by crafting logp, ref_logp with known values. Output is elementwise (same shape as input) , confirm NO reduction (or determine if reduced).\n- reverse_kl(logp, ref_logp): determine exact formula. Likely exp(logp-ref_logp)*(logp-ref_logp) or the k3-style, or simply (logp - ref_logp). Probe with known values to pin down. Elementwise?\n- symmetric_kl(logp, ref_logp): symmetric combination. Determine formula (e.g. reverse_kl + forward_kl, or 0.5*(...)). Probe.\n- importance_ratio(logp, old_logp, clip): ratio = exp(logp - old_logp). clip is a scalar; if clip is not None the ratio is clamped to [1-clip, 1+clip] OR [clip stuff] , determine exact clamp bounds and whether clip=None means no clamp. Probe with clip value e.g. 0.2 and inputs that exceed bounds.\n- clip_fraction(logp, old_logp, clip): fraction of elements where the importance ratio would be clipped, i.e. mean over elements of (|ratio-1| > clip) or (ratio outside [1-clip,1+clip]). Determine exact condition (strict > vs >=) and reduction (mean, scalar). Probe with values exactly at boundary.\n- bradley_terry_logit(chosen_reward, rejected_reward, beta): likely beta*(chosen_reward - rejected_reward), or logsigmoid thereof, or -logsigmoid. Probe: is it just the logit beta*(rc-rr), or a loss -logsigmoid(beta*(rc-rr))? The name says 'logit' so probably beta*(chosen-rejected). Confirm by probing.\n\nVERIFICATION: Write a test comparing impl to reforacle over >=30 random trials each, diverse shapes, all estimator strings, clip=None and clip values, boundary cases. torch.allclose(atol=1e-6,rtol=1e-5,equal_nan=True). Iterate until all pass.\n\nReturn ONLY final Python source for these 6 functions (no prose/markdown), standalone torch-only (`import torch` assumed). Helpers nested or `_kl_` prefixed. No reforacle reference in returned code. Also, in your final message, ALONGSIDE the code, include a short note stating the exact estimator strings supported and the exact sign convention you found (I need this for composite objectives).",
  "subagent_type": "general-purpose"
}

, output, 
{"type": "text", "text": "All 60 trials pass across scalar and vector shapes, diverse beta, and label_smoothing (including 0).\n\n```python\ndef dpo_loss(pc, pr, rc, rr, beta, label_smoothing):\n    import torch\n    logits = beta * ((pc - rc) - (pr - rr))\n    loss = (\n        -(1 - label_smoothing) * torch.nn.functional.logsigmoid(logits)\n        - label_smoothing * torch.nn.functional.logsigmoid(-logits)\n    )\n    return loss.mean()\n\n\ndef ipo_loss(pc, pr, rc, rr, beta):\n    h = (pc - rc) - (pr - rr)\n    loss = (h - 1 / (2 * beta)) ** 2\n    return loss.mean()\n```\n\nKey facts for composite objectives:\n\n- DPO logits formula: `logits = beta * ((pc - rc) - (pr - rr))`\n- DPO label_smoothing formula (per-example, before reduction): `-(1 - label_smoothing) * logsigmoid(logits) - label_smoothing * logsigmoid(-logits)`. At `label_smoothing=0` this collapses to `-logsigmoid(logits)`.\n- IPO formula (per-example): `((pc - rc) - (pr - rr) - 1/(2*beta))**2`, i.e. `(h - 1/(2*beta))**2` with `h = (pc-rc)-(pr-rr)`. Beta appears only inside the `1/(2*beta)` margin term, not as a multiplier on `h`.\n- Reduction convention: BOTH functions return a MEAN scalar (shape `[]`), not per-example. Verified by passing N=4 vectors and getting a scalar equal to the per-example `.mean()`.\n\nVerified over 60 trials (scalar and `(N,)` shapes, beta in ~[0.05, 2.05], label_smoothing including 0) with `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)` , all pass."}

{"type": "text", "text": "agentId: ab14fc7203cb72eaf (use SendMessage with to: 'ab14fc7203cb72eaf', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 12305\ntool_uses: 3\nduration_ms: 37779</usage>"}

{"type": "text", "text": "All 60 trials plus edge cases pass with zero failures. Here is the final implementation.\n\n```python\nimport torch\n\n\ndef _lp_masked_mean(values, mask):\n    denom = mask.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return (values * mask).sum() / denom\n\n\ndef logsumexp(x, axis):\n    return torch.logsumexp(x, dim=axis)\n\n\ndef log_softmax(x, axis):\n    return x - torch.logsumexp(x, dim=axis, keepdim=True)\n\n\ndef token_logprobs(logits, labels):\n    lsm = torch.log_softmax(logits, dim=-1)\n    return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n\n\ndef selective_logprobs(logits, labels, mask):\n    lp = token_logprobs(logits, labels)\n    return (lp * mask).sum(dim=-1)\n\n\ndef sequence_logprob(logits, labels, mask, length_normalize):\n    lp = token_logprobs(logits, labels)\n    total = (lp * mask).sum(dim=-1)\n    if length_normalize:\n        denom = mask.sum(dim=-1)\n        safe = torch.where(denom > 0, denom, torch.ones_like(denom))\n        out = total / safe\n        return torch.where(denom > 0, out, torch.zeros_like(out))\n    return total\n\n\ndef logprob_at_temperature(logits, labels, temperature):\n    return token_logprobs(logits / temperature, labels)\n\n\ndef cross_entropy(logits, labels, ignore_index):\n    labels = labels.long()\n    flat_logits = logits.reshape(-1, logits.shape[-1])\n    flat_labels = labels.reshape(-1)\n    lsm = torch.log_softmax(flat_logits, dim=-1)\n    valid = (flat_labels != ignore_index)\n    safe_labels = torch.where(valid, flat_labels, torch.zeros_like(flat_labels))\n    nll = -lsm.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)\n    valid_f = valid.to(nll.dtype)\n    denom = valid_f.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=nll.dtype, device=nll.device)\n    return (nll * valid_f).sum() / denom\n\n\ndef smoothed_nll(logits, labels, smoothing):\n    lsm = torch.log_softmax(logits, dim=-1)\n    nll = -lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n    smooth = -lsm.mean(dim=-1)\n    loss = (1.0 - smoothing) * nll + smoothing * smooth\n    return loss.mean()\n\n\ndef entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    return _lp_masked_mean(ent, mask)\n\n\ndef normalized_entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    vocab = logits.shape[-1]\n    norm = ent / torch.log(torch.tensor(float(vocab), dtype=ent.dtype, device=ent.device))\n    return _lp_masked_mean(norm, mask)\n```\n\nRecovered conventions (all verified over 60 diverse trials + edge cases, 0 failures at atol=1e-6/rtol=1e-5):\n- **logsumexp**: standard, no keepdim; **log_softmax**: `x - logsumexp(x, keepdim=True)`.\n- **token_logprobs**: log_softmax over last dim, gather labels; output shape = labels shape.\n- **selective_logprobs**: `sum(token_logprobs * mask, dim=-1)` , sums over the sequence (last) axis, dropping it.\n- **sequence_logprob**: same masked sum; `length_normalize=True` divides by `mask.sum(-1)`, returning 0 when the mask is all-zero (guarded div).\n- **logprob_at_temperature**: divides logits by temperature before log_softmax.\n- **cross_entropy**: flattens to (-1, V), mean of NLL over non-ignored positions (`sum/count`); returns 0 (not NaN) when all positions are ignored. Works for any ignore_index value.\n- **smoothed_nll**: `mean[(1-s)*nll + s*(-mean_vocab log_softmax)]`.\n- **entropy**: per-position natural-log entropy `-sum(p*log p)`, masked-mean over positions.\n- **normalized_entropy**: entropy / log(vocab_size), masked-mean; returns 0 on empty mask.\n\nNote: `masked_mean` on an all-zero mask returns 0, which is the guard reused by entropy/normalized_entropy."}

{"type": "text", "text": "agentId: acac49472241baa47 (use SendMessage with to: 'acac49472241baa47', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 26571\ntool_uses: 19\nduration_ms: 171626</usage>"}

{"type": "text", "text": "All 1762 tests pass across diverse shapes, group sizes, gamma/lam values, scale_by_std both ways, clip values, and boundary cases.\n\n```python\nimport torch\n\n\ndef _rl_masked_mean(x, mask):\n    return (x * mask).sum() / mask.sum()\n\n\ndef grpo_advantages(rewards, group_size, scale_by_std):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    if scale_by_std:\n        std = groups.std(dim=1, unbiased=False, keepdim=True)\n        adv = adv / (std + 1e-6)\n    return adv.reshape(-1)\n\n\ndef rloo_advantages(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    group_sum = groups.sum(dim=1, keepdim=True)\n    baseline = (group_sum - groups) / (group_size - 1)\n    adv = groups - baseline\n    return adv.reshape(-1)\n\n\ndef group_mean_baseline(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    return adv.reshape(-1)\n\n\ndef gae(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    adv = torch.zeros_like(rewards)\n    gae_acc = torch.zeros_like(rewards[0])\n    next_v = next_value\n    for t in range(T - 1, -1, -1):\n        delta = rewards[t] + gamma * next_v - values[t]\n        gae_acc = delta + gamma * lam * gae_acc\n        adv[t] = gae_acc\n        next_v = values[t]\n    return adv\n\n\ndef discounted_returns(rewards, gamma):\n    T = rewards.shape[0]\n    out = torch.zeros_like(rewards)\n    acc = torch.zeros_like(rewards[0])\n    for t in range(T - 1, -1, -1):\n        acc = rewards[t] + gamma * acc\n        out[t] = acc\n    return out\n\n\ndef lambda_returns(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    out = torch.zeros_like(rewards)\n    next_ret = next_value\n    for t in range(T - 1, -1, -1):\n        next_v = next_value if t == T - 1 else values[t + 1]\n        next_ret = rewards[t] + gamma * ((1 - lam) * next_v + lam * next_ret)\n        out[t] = next_ret\n    return out\n\n\ndef clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):\n    ratio = torch.exp(logp - old_logp)\n    unclipped = ratio * advantages\n    clipped = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) * advantages\n    per = -torch.min(unclipped, clipped)\n    return _rl_masked_mean(per, mask)\n\n\ndef value_loss(values, old_values, returns, clip):\n    unclipped = (values - returns) ** 2\n    vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)\n    clipped = (vpred_clipped - returns) ** 2\n    return 0.5 * torch.max(unclipped, clipped).mean()\n\n\ndef huber_value_loss(values, returns, delta):\n    x = values - returns\n    absx = x.abs()\n    huber = torch.where(absx <= delta, 0.5 * x ** 2, delta * (absx - 0.5 * delta))\n    return huber.mean()\n```\n\nKey conventions you asked about for composite objectives:\n\n- **GAE / returns time-axis & shape**: Time is the **first axis (dim 0)**. Inputs may be 1-D `(T,)` or N-D `(T, ...)` (e.g. `(T, B)`); the recursion runs over dim 0 and per-timestep slices broadcast over trailing dims. `next_value` may be a scalar or shape `(...)` matching the trailing dims. `gae` returns **advantages only** (a single tensor, not a tuple), same shape as `rewards`. `lambda_returns` returns TD(λ) returns and exactly equals `gae + values`. `discounted_returns` also reduces along dim 0.\n\n- **grpo_advantages / group_mean_baseline shape**: Input is flattened then reshaped to `(-1, group_size)` (row-major grouping), and **output is always flat `(N,)`** regardless of input shape. std is **biased** (`unbiased=False`), eps = **1e-6**. Note `group_mean_baseline` returns the **advantages** (`reward - group_mean`), not the raw baseline.\n\n- **clipped_pg_loss sign & clamp**: `ratio = exp(logp - old_logp)`, `clamp(ratio, 1 - clip_low, 1 + clip_high)` (clip_low is the lower bound, clip_high the upper), then **`-min(unclipped, clipped)`** (negative , it is a loss), reduced by **masked mean = `sum(per*mask)/sum(mask)`**.\n\n- **value_loss 0.5 factor**: Yes , `0.5 * mean(max(unclipped, clipped))`, using **max** of clipped vs unclipped squared errors, plain `.mean()` over all elements. `clip` is always numeric (the oracle does not special-case `None`; passing `None` errors). `clip=0.0` degenerates to `0.5*mean((old_values - returns)^2)` vs unclipped max.\n\n- **huber_value_loss**: standard smooth-L1 (`0.5*x^2` for `|x|<=delta`, else `delta*(|x|-0.5*delta)`), reduced by plain mean; equals `F.smooth_l1_loss(..., beta=delta)`."}

{"type": "text", "text": "agentId: ae12d609b60dce0b3 (use SendMessage with to: 'ae12d609b60dce0b3', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 25299\ntool_uses: 11\nduration_ms: 165387</usage>"}

{"type": "text", "text": "All edge cases pass. Implementation complete and verified in `/app/objectives.py`.\n\nFinal Python source for the 6 functions:\n\n```python\nimport torch\n\n\ndef kl_penalty(logp, ref_logp, estimator):\n    logratio = logp - ref_logp\n    if estimator == 'k1':\n        return logratio\n    elif estimator == 'k2':\n        return 0.5 * logratio * logratio\n    elif estimator == 'k3':\n        return torch.exp(-logratio) + logratio - 1\n    else:\n        raise ValueError(estimator)\n\n\ndef reverse_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    return torch.exp(-logratio) + logratio - 1\n\n\ndef symmetric_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    reverse = torch.exp(-logratio) + logratio - 1\n    forward = torch.exp(logratio) - logratio - 1\n    return 0.5 * (reverse + forward)\n\n\ndef importance_ratio(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    if clip is not None:\n        ratio = torch.clamp(ratio, 1 - clip, 1 + clip)\n    return ratio\n\n\ndef clip_fraction(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    return ((ratio > 1 + clip) | (ratio < 1 - clip)).float().mean()\n\n\ndef bradley_terry_logit(chosen_reward, rejected_reward, beta):\n    return beta * (chosen_reward - rejected_reward)\n```\n\nNote on conventions (for composite objectives):\n\n- Supported estimator strings for `kl_penalty`: exactly `'k1'`, `'k2'`, `'k3'`. Any other string raises `ValueError(estimator)`.\n- Sign convention: `logratio = logp - ref_logp` (policy minus reference). Then:\n  - `k1 = logratio` (elementwise, no reduction, same shape as input)\n  - `k2 = 0.5 * logratio**2`\n  - `k3 = exp(-logratio) + logratio - 1` (the non-negative unbiased estimator, using `-logratio` inside exp)\n- `reverse_kl` equals the k3 form: `exp(-logratio) + logratio - 1`. Elementwise, no reduction.\n- `symmetric_kl = 0.5 * (reverse + forward)` where `reverse = exp(-logratio)+logratio-1` and `forward = exp(logratio)-logratio-1`. Elementwise.\n- `importance_ratio`: `ratio = exp(logp - old_logp)`; `clip=None` means no clamp; otherwise clamp to `[1-clip, 1+clip]`. Elementwise.\n- `clip_fraction`: scalar (0-dim tensor) mean of the strict out-of-range condition `(ratio > 1+clip) | (ratio < 1-clip)`. Important: it is NOT `|ratio-1| > clip` , that form disagrees on floating-point boundary values (e.g. `exp(log(1.2))` = 1.20000004 counts as clipped under abs-form but not under the range-form the oracle uses).\n- `bradley_terry_logit = beta * (chosen_reward - rejected_reward)` , a plain logit, no sigmoid/logsigmoid. Elementwise.\n\nVerification: 452 comparisons across 30+ random trials with diverse shapes `(4,), (3,5), (2,3,4), (1,), (10,), (7,2)`, all three estimators, `clip=None` and clip values `{0.0,0.05,0.1,0.2,0.3,0.5,1.0}`, boundary ratios exactly at `1±clip`, and inf/nan inputs , all pass `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)`."}

{"type": "text", "text": "agentId: a5acd56804a9a0d15 (use SendMessage with to: 'a5acd56804a9a0d15', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 29330\ntool_uses: 33\nduration_ms: 212902</usage>"}

{"type": "text", "text": "All 10818 randomized trials pass across diverse shapes (1D/2D/3D), mask kinds (full/zero/sparse/half), all axes, both shift modes, multiple eps values, and empty masks. Combined with the earlier suites and dtype tests, the implementation exactly reproduces `reforacle`.\n\nRecovered conventions:\n- **masked_mean**: `sum(values*mask)/sum(mask)` (mask = float weights), safe-divides to 0 on empty mask; respects `axis` with the reduced dim dropped.\n- **masked_sum**: `sum(values*mask)` over axis (or global).\n- **masked_max**: global scalar, excludes masked positions via `where(mask, values, -inf)`; empty mask returns 0 in the input dtype.\n- **whiten**: masked biased mean/var (÷sum(mask)), `(x-mean)/sqrt(var+1e-8)` applied to all elements; if `shift_mean` is False, adds mean back. No output masking.\n- **masked_whiten**: same as whiten but multiplies the whitened result by mask before optionally adding mean back.\n- **advantage_mean_std**: returns `stack([masked_mean, sqrt(masked_biased_var)])` , biased std, no eps, empty→[0,0].\n- **normalize**: global `(x - x.mean()) / (x.std(unbiased=True) + eps)` , unbiased std, eps added to std (not inside sqrt).\n\nFinal implementation:\n\n```python\nimport torch\n\ndef masked_mean(values, mask, axis=None):\n    if axis is None:\n        s = (values * mask).sum()\n        d = mask.sum()\n    else:\n        s = (values * mask).sum(dim=axis)\n        d = mask.sum(dim=axis)\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef masked_sum(values, mask, axis=None):\n    if axis is None:\n        return (values * mask).sum()\n    return (values * mask).sum(dim=axis)\n\ndef masked_max(values, mask):\n    neg_inf = torch.finfo(values.dtype).min\n    masked_vals = torch.where(mask.bool(), values, torch.full_like(values, neg_inf))\n    if mask.sum() == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return masked_vals.max()\n\ndef _maskprim_safe_mean(values, mask):\n    s = (values * mask).sum()\n    d = mask.sum()\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef _maskprim_mvar(values, mask):\n    mean = _maskprim_safe_mean(values, mask)\n    centered = values - mean\n    var = _maskprim_safe_mean(centered * centered, mask)\n    return mean, var\n\ndef whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef masked_whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    whitened = whitened * mask\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef advantage_mean_std(advantages, mask):\n    mean, var = _maskprim_mvar(advantages, mask)\n    std = torch.sqrt(var)\n    return torch.stack([mean, std])\n\ndef normalize(x, eps):\n    mean = x.mean()\n    std = x.std(unbiased=True)\n    return (x - mean) / (std + eps)\n```"}

{"type": "text", "text": "agentId: af76df31e94c64aa5 (use SendMessage with to: 'af76df31e94c64aa5', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 33906\ntool_uses: 19\nduration_ms: 267714</usage>"}

{"type": "text", "text": "All final edge cases pass, including float64 dtype preservation. The implementation is verified. Here is the final standalone source.\n\n```python\nimport torch\n\n\ndef top_p_mask(probs, p):\n    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)\n    cumsum = torch.cumsum(sorted_probs, dim=-1)\n    cumsum_before = cumsum - sorted_probs\n    keep = cumsum_before < p\n    keep[..., 0] = True\n    kept = sorted_probs * keep.to(sorted_probs.dtype)\n    kept = kept / kept.sum(dim=-1, keepdim=True)\n    out = torch.zeros_like(kept)\n    out.scatter_(-1, sorted_idx, kept)\n    return out\n\n\ndef top_k_mask(logits, k):\n    kth = torch.topk(logits, k, dim=-1).values[..., -1:]\n    return logits >= kth\n\n\ndef argmax_tokens(logits):\n    flipped = torch.flip(logits, dims=[-1])\n    n = logits.shape[-1]\n    idx = torch.argmax(flipped, dim=-1)\n    return (n - 1) - idx\n\n\ndef mode_label(labels):\n    def _mode_1d(v):\n        uniq, counts = torch.unique(v, return_counts=True)\n        max_count = counts.max()\n        candidates = uniq[counts == max_count]\n        return candidates.max()\n\n    if labels.dim() == 1:\n        return _mode_1d(labels)\n    return torch.stack([_mode_1d(labels[i]) for i in range(labels.shape[0])])\n\n\ndef _positional_quantile(x, q, interpolation):\n    xs, _ = torch.sort(x, dim=-1)\n    n = xs.shape[0]\n    pos = (n - 1) * q\n    if interpolation == \"lower\":\n        return xs[int(pos // 1)]\n    lo = int(pos // 1)\n    hi = min(lo + 1, n - 1)\n    frac = pos - lo\n    return xs[lo] * (1 - frac) + xs[hi] * frac\n\n\ndef median_reward(rewards):\n    return _positional_quantile(rewards, 0.5, \"linear\")\n\n\ndef quantile_lower(x, q):\n    return _positional_quantile(x, q, \"lower\")\n\n\ndef pad_mask_from_lengths(lengths, max_len):\n    arange = torch.arange(max_len, device=lengths.device).unsqueeze(0)\n    return (arange < lengths.unsqueeze(1)).to(torch.float32)\n\n\ndef first_nonzero_index(mask):\n    nz = mask != 0\n    any_nz = nz.any(dim=-1)\n    idx = torch.argmax(nz.to(torch.long), dim=-1)\n    return torch.where(any_nz, idx, torch.full_like(idx, -1))\n\n\ndef cumulative_max(x):\n    out = torch.empty_like(x)\n    out[0] = x[0]\n    if x.shape[0] > 1:\n        out[1:] = torch.cummax(x, dim=0)[0][:-1]\n    return out\n\n\ndef bucketize_reward(rewards, boundaries):\n    return torch.bucketize(rewards, boundaries, right=True)\n```\n\n## Summary\n\nI reverse-engineered all 10 functions by probing `reforacle` and verified against 100+ diverse trials (including heavy tie/edge cases) with zero mismatches. Key recovered conventions, several of which are non-obvious:\n\n- **top_p_mask**: nucleus over last dim; keep where `cumsum_before < p` (HF shift, crossing token included), always keeps ≥1 token, output is **renormalized probabilities** in original order (not a boolean mask), zeros elsewhere.\n- **top_k_mask**: **threshold-based** (`logit >= kth-largest value`), so ties at the threshold are all kept and the count may exceed k. Returns bool, last dim.\n- **argmax_tokens**: last dim, but returns the **LAST** maximum index on ties (not torch.argmax's first).\n- **mode_label**: most frequent value per row (dim 0 for 1D, per-row for 2D); tie-break returns the **largest** value (not torch.mode which returns smallest). Handles negatives.\n- **median_reward / quantile_lower**: reduce along dim 0, but with a non-standard rule , **sort along the last dim, then take the positional quantile across dim 0 without re-sorting** (median=linear interp, quantile_lower=floor/lower). Standard `torch.quantile(dim=0)` does NOT reproduce this because it re-sorts each column. For 1D this coincides with normal quantile.\n- **pad_mask_from_lengths**: `(N, max_len)` float32, `arange < length`.\n- **first_nonzero_index**: per-row along last dim, returns **-1** when a row is all zero.\n- **cumulative_max**: exclusive right-shifted running max along **dim 0** (`out[0]=x[0]`, `out[i]=max(x[0..i-1])`).\n- **bucketize_reward**: `torch.bucketize(..., right=True)`.\n\nFiles: `/app/myimpl.py` (final implementation), verification scripts `/app/verify2.py`, `/app/verify4.py`, `/app/verify5.py`. The `/app/objectives.py` scaffold was left unmodified (it contains many other stubs unrelated to this task)."}

{"type": "text", "text": "agentId: aca416a9a99d1536f (use SendMessage with to: 'aca416a9a99d1536f', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 63559\ntool_uses: 37\nduration_ms: 552401</usage>"}

{
  "description": "Reverse-engineer masking primitives",
  "prompt": "You are reverse-engineering a sealed reference module `reforacle` (importable via `import reforacle`, source unreadable). You must implement Python functions that reproduce `reforacle`'s exact numerical behavior. Your ONLY dependency may be `torch`. You recover behavior ONLY by calling reforacle and comparing , never decompile.\n\nRun `python3 /app/probe.py` to see the calling pattern. The oracle is called like `reforacle.masked_mean(values, mask)`.\n\nImplement EXACTLY these functions (keep signatures identical):\n\ndef masked_mean(values, mask, axis=None):\ndef masked_sum(values, mask, axis=None):\ndef masked_max(values, mask):\ndef whiten(values, mask, shift_mean):\ndef masked_whiten(values, mask, shift_mean):\ndef advantage_mean_std(advantages, mask):\ndef normalize(x, eps):\n\nCRITICAL: Probe each function extensively before implementing. Conventions are non-obvious and you WILL be graded on held-out inputs, so you must nail the EXACT convention, not just match a few cases. Things to determine by probing:\n- masked_mean/masked_sum: behavior when axis=None (global) vs axis given (e.g. axis=0, axis=1, negative axis). Shape of output. Behavior with empty mask (returns 0.0). Whether mask is treated as float weights or boolean. Test 1D, 2D, 3D tensors. Does masked_mean divide by sum(mask) or by count? What if partial mask on 2D with axis?\n- masked_max: what is returned for masked-out positions? global or per-axis? what for empty mask?\n- whiten vs masked_whiten: (x-mean)/std. Determine: does std use unbiased (N-1) or biased (N) correction? Is there an epsilon added to std/variance (probe with known values to find eps, e.g. 1e-8)? shift_mean=True/False controls whether mean is subtracted back / added. Probe both True and False. whiten likely unmasked (uses all values) while masked_whiten uses mask. For masked_whiten determine masked mean/var convention (divide by sum(mask) or sum(mask)-1?).\n- advantage_mean_std: returns what? Probably a tuple (mean, std) or normalized advantages. Probe return type/shape. Masked. unbiased or biased std? eps?\n- normalize(x, eps): (x-mean)/(std+eps) or (x-mean)/sqrt(var+eps)? global? unbiased std? Determine exact eps usage by probing with crafted inputs.\n\nTo find eps and correction: use small tensors where you can compute candidates by hand, e.g. x=[0.,1.] gives mean 0.5, biased var 0.25 (std 0.5), unbiased var 0.5 (std 0.707). Compare reforacle output to candidates.\n\nVERIFICATION REQUIREMENT: Write a test script that compares your implementation against reforacle over MANY random trials (>=30 each) with diverse shapes (1D/2D/3D), diverse masks (full, partial, single-element, all-zero), and diverse dtypes. Use torch.allclose(a,b,atol=1e-6,rtol=1e-5) (equal_nan=True). Only trust an implementation that passes ALL trials including edge cases. Iterate until perfect.\n\nReturn ONLY the final Python source code for these 7 functions (no prose, no markdown fences), as standalone functions depending only on torch (module-level `import torch` assumed present). If you need helpers, define them as nested functions inside each function or with a unique prefix like `_maskprim_`. Do NOT reference reforacle in returned code.",
  "subagent_type": "general-purpose"
}

, output, 
{"type": "text", "text": "All 60 trials pass across scalar and vector shapes, diverse beta, and label_smoothing (including 0).\n\n```python\ndef dpo_loss(pc, pr, rc, rr, beta, label_smoothing):\n    import torch\n    logits = beta * ((pc - rc) - (pr - rr))\n    loss = (\n        -(1 - label_smoothing) * torch.nn.functional.logsigmoid(logits)\n        - label_smoothing * torch.nn.functional.logsigmoid(-logits)\n    )\n    return loss.mean()\n\n\ndef ipo_loss(pc, pr, rc, rr, beta):\n    h = (pc - rc) - (pr - rr)\n    loss = (h - 1 / (2 * beta)) ** 2\n    return loss.mean()\n```\n\nKey facts for composite objectives:\n\n- DPO logits formula: `logits = beta * ((pc - rc) - (pr - rr))`\n- DPO label_smoothing formula (per-example, before reduction): `-(1 - label_smoothing) * logsigmoid(logits) - label_smoothing * logsigmoid(-logits)`. At `label_smoothing=0` this collapses to `-logsigmoid(logits)`.\n- IPO formula (per-example): `((pc - rc) - (pr - rr) - 1/(2*beta))**2`, i.e. `(h - 1/(2*beta))**2` with `h = (pc-rc)-(pr-rr)`. Beta appears only inside the `1/(2*beta)` margin term, not as a multiplier on `h`.\n- Reduction convention: BOTH functions return a MEAN scalar (shape `[]`), not per-example. Verified by passing N=4 vectors and getting a scalar equal to the per-example `.mean()`.\n\nVerified over 60 trials (scalar and `(N,)` shapes, beta in ~[0.05, 2.05], label_smoothing including 0) with `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)` , all pass."}

{"type": "text", "text": "agentId: ab14fc7203cb72eaf (use SendMessage with to: 'ab14fc7203cb72eaf', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 12305\ntool_uses: 3\nduration_ms: 37779</usage>"}

{"type": "text", "text": "All 60 trials plus edge cases pass with zero failures. Here is the final implementation.\n\n```python\nimport torch\n\n\ndef _lp_masked_mean(values, mask):\n    denom = mask.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return (values * mask).sum() / denom\n\n\ndef logsumexp(x, axis):\n    return torch.logsumexp(x, dim=axis)\n\n\ndef log_softmax(x, axis):\n    return x - torch.logsumexp(x, dim=axis, keepdim=True)\n\n\ndef token_logprobs(logits, labels):\n    lsm = torch.log_softmax(logits, dim=-1)\n    return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n\n\ndef selective_logprobs(logits, labels, mask):\n    lp = token_logprobs(logits, labels)\n    return (lp * mask).sum(dim=-1)\n\n\ndef sequence_logprob(logits, labels, mask, length_normalize):\n    lp = token_logprobs(logits, labels)\n    total = (lp * mask).sum(dim=-1)\n    if length_normalize:\n        denom = mask.sum(dim=-1)\n        safe = torch.where(denom > 0, denom, torch.ones_like(denom))\n        out = total / safe\n        return torch.where(denom > 0, out, torch.zeros_like(out))\n    return total\n\n\ndef logprob_at_temperature(logits, labels, temperature):\n    return token_logprobs(logits / temperature, labels)\n\n\ndef cross_entropy(logits, labels, ignore_index):\n    labels = labels.long()\n    flat_logits = logits.reshape(-1, logits.shape[-1])\n    flat_labels = labels.reshape(-1)\n    lsm = torch.log_softmax(flat_logits, dim=-1)\n    valid = (flat_labels != ignore_index)\n    safe_labels = torch.where(valid, flat_labels, torch.zeros_like(flat_labels))\n    nll = -lsm.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)\n    valid_f = valid.to(nll.dtype)\n    denom = valid_f.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=nll.dtype, device=nll.device)\n    return (nll * valid_f).sum() / denom\n\n\ndef smoothed_nll(logits, labels, smoothing):\n    lsm = torch.log_softmax(logits, dim=-1)\n    nll = -lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n    smooth = -lsm.mean(dim=-1)\n    loss = (1.0 - smoothing) * nll + smoothing * smooth\n    return loss.mean()\n\n\ndef entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    return _lp_masked_mean(ent, mask)\n\n\ndef normalized_entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    vocab = logits.shape[-1]\n    norm = ent / torch.log(torch.tensor(float(vocab), dtype=ent.dtype, device=ent.device))\n    return _lp_masked_mean(norm, mask)\n```\n\nRecovered conventions (all verified over 60 diverse trials + edge cases, 0 failures at atol=1e-6/rtol=1e-5):\n- **logsumexp**: standard, no keepdim; **log_softmax**: `x - logsumexp(x, keepdim=True)`.\n- **token_logprobs**: log_softmax over last dim, gather labels; output shape = labels shape.\n- **selective_logprobs**: `sum(token_logprobs * mask, dim=-1)` , sums over the sequence (last) axis, dropping it.\n- **sequence_logprob**: same masked sum; `length_normalize=True` divides by `mask.sum(-1)`, returning 0 when the mask is all-zero (guarded div).\n- **logprob_at_temperature**: divides logits by temperature before log_softmax.\n- **cross_entropy**: flattens to (-1, V), mean of NLL over non-ignored positions (`sum/count`); returns 0 (not NaN) when all positions are ignored. Works for any ignore_index value.\n- **smoothed_nll**: `mean[(1-s)*nll + s*(-mean_vocab log_softmax)]`.\n- **entropy**: per-position natural-log entropy `-sum(p*log p)`, masked-mean over positions.\n- **normalized_entropy**: entropy / log(vocab_size), masked-mean; returns 0 on empty mask.\n\nNote: `masked_mean` on an all-zero mask returns 0, which is the guard reused by entropy/normalized_entropy."}

{"type": "text", "text": "agentId: acac49472241baa47 (use SendMessage with to: 'acac49472241baa47', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 26571\ntool_uses: 19\nduration_ms: 171626</usage>"}

{"type": "text", "text": "All 1762 tests pass across diverse shapes, group sizes, gamma/lam values, scale_by_std both ways, clip values, and boundary cases.\n\n```python\nimport torch\n\n\ndef _rl_masked_mean(x, mask):\n    return (x * mask).sum() / mask.sum()\n\n\ndef grpo_advantages(rewards, group_size, scale_by_std):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    if scale_by_std:\n        std = groups.std(dim=1, unbiased=False, keepdim=True)\n        adv = adv / (std + 1e-6)\n    return adv.reshape(-1)\n\n\ndef rloo_advantages(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    group_sum = groups.sum(dim=1, keepdim=True)\n    baseline = (group_sum - groups) / (group_size - 1)\n    adv = groups - baseline\n    return adv.reshape(-1)\n\n\ndef group_mean_baseline(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    return adv.reshape(-1)\n\n\ndef gae(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    adv = torch.zeros_like(rewards)\n    gae_acc = torch.zeros_like(rewards[0])\n    next_v = next_value\n    for t in range(T - 1, -1, -1):\n        delta = rewards[t] + gamma * next_v - values[t]\n        gae_acc = delta + gamma * lam * gae_acc\n        adv[t] = gae_acc\n        next_v = values[t]\n    return adv\n\n\ndef discounted_returns(rewards, gamma):\n    T = rewards.shape[0]\n    out = torch.zeros_like(rewards)\n    acc = torch.zeros_like(rewards[0])\n    for t in range(T - 1, -1, -1):\n        acc = rewards[t] + gamma * acc\n        out[t] = acc\n    return out\n\n\ndef lambda_returns(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    out = torch.zeros_like(rewards)\n    next_ret = next_value\n    for t in range(T - 1, -1, -1):\n        next_v = next_value if t == T - 1 else values[t + 1]\n        next_ret = rewards[t] + gamma * ((1 - lam) * next_v + lam * next_ret)\n        out[t] = next_ret\n    return out\n\n\ndef clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):\n    ratio = torch.exp(logp - old_logp)\n    unclipped = ratio * advantages\n    clipped = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) * advantages\n    per = -torch.min(unclipped, clipped)\n    return _rl_masked_mean(per, mask)\n\n\ndef value_loss(values, old_values, returns, clip):\n    unclipped = (values - returns) ** 2\n    vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)\n    clipped = (vpred_clipped - returns) ** 2\n    return 0.5 * torch.max(unclipped, clipped).mean()\n\n\ndef huber_value_loss(values, returns, delta):\n    x = values - returns\n    absx = x.abs()\n    huber = torch.where(absx <= delta, 0.5 * x ** 2, delta * (absx - 0.5 * delta))\n    return huber.mean()\n```\n\nKey conventions you asked about for composite objectives:\n\n- **GAE / returns time-axis & shape**: Time is the **first axis (dim 0)**. Inputs may be 1-D `(T,)` or N-D `(T, ...)` (e.g. `(T, B)`); the recursion runs over dim 0 and per-timestep slices broadcast over trailing dims. `next_value` may be a scalar or shape `(...)` matching the trailing dims. `gae` returns **advantages only** (a single tensor, not a tuple), same shape as `rewards`. `lambda_returns` returns TD(λ) returns and exactly equals `gae + values`. `discounted_returns` also reduces along dim 0.\n\n- **grpo_advantages / group_mean_baseline shape**: Input is flattened then reshaped to `(-1, group_size)` (row-major grouping), and **output is always flat `(N,)`** regardless of input shape. std is **biased** (`unbiased=False`), eps = **1e-6**. Note `group_mean_baseline` returns the **advantages** (`reward - group_mean`), not the raw baseline.\n\n- **clipped_pg_loss sign & clamp**: `ratio = exp(logp - old_logp)`, `clamp(ratio, 1 - clip_low, 1 + clip_high)` (clip_low is the lower bound, clip_high the upper), then **`-min(unclipped, clipped)`** (negative , it is a loss), reduced by **masked mean = `sum(per*mask)/sum(mask)`**.\n\n- **value_loss 0.5 factor**: Yes , `0.5 * mean(max(unclipped, clipped))`, using **max** of clipped vs unclipped squared errors, plain `.mean()` over all elements. `clip` is always numeric (the oracle does not special-case `None`; passing `None` errors). `clip=0.0` degenerates to `0.5*mean((old_values - returns)^2)` vs unclipped max.\n\n- **huber_value_loss**: standard smooth-L1 (`0.5*x^2` for `|x|<=delta`, else `delta*(|x|-0.5*delta)`), reduced by plain mean; equals `F.smooth_l1_loss(..., beta=delta)`."}

{"type": "text", "text": "agentId: ae12d609b60dce0b3 (use SendMessage with to: 'ae12d609b60dce0b3', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 25299\ntool_uses: 11\nduration_ms: 165387</usage>"}

{"type": "text", "text": "All edge cases pass. Implementation complete and verified in `/app/objectives.py`.\n\nFinal Python source for the 6 functions:\n\n```python\nimport torch\n\n\ndef kl_penalty(logp, ref_logp, estimator):\n    logratio = logp - ref_logp\n    if estimator == 'k1':\n        return logratio\n    elif estimator == 'k2':\n        return 0.5 * logratio * logratio\n    elif estimator == 'k3':\n        return torch.exp(-logratio) + logratio - 1\n    else:\n        raise ValueError(estimator)\n\n\ndef reverse_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    return torch.exp(-logratio) + logratio - 1\n\n\ndef symmetric_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    reverse = torch.exp(-logratio) + logratio - 1\n    forward = torch.exp(logratio) - logratio - 1\n    return 0.5 * (reverse + forward)\n\n\ndef importance_ratio(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    if clip is not None:\n        ratio = torch.clamp(ratio, 1 - clip, 1 + clip)\n    return ratio\n\n\ndef clip_fraction(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    return ((ratio > 1 + clip) | (ratio < 1 - clip)).float().mean()\n\n\ndef bradley_terry_logit(chosen_reward, rejected_reward, beta):\n    return beta * (chosen_reward - rejected_reward)\n```\n\nNote on conventions (for composite objectives):\n\n- Supported estimator strings for `kl_penalty`: exactly `'k1'`, `'k2'`, `'k3'`. Any other string raises `ValueError(estimator)`.\n- Sign convention: `logratio = logp - ref_logp` (policy minus reference). Then:\n  - `k1 = logratio` (elementwise, no reduction, same shape as input)\n  - `k2 = 0.5 * logratio**2`\n  - `k3 = exp(-logratio) + logratio - 1` (the non-negative unbiased estimator, using `-logratio` inside exp)\n- `reverse_kl` equals the k3 form: `exp(-logratio) + logratio - 1`. Elementwise, no reduction.\n- `symmetric_kl = 0.5 * (reverse + forward)` where `reverse = exp(-logratio)+logratio-1` and `forward = exp(logratio)-logratio-1`. Elementwise.\n- `importance_ratio`: `ratio = exp(logp - old_logp)`; `clip=None` means no clamp; otherwise clamp to `[1-clip, 1+clip]`. Elementwise.\n- `clip_fraction`: scalar (0-dim tensor) mean of the strict out-of-range condition `(ratio > 1+clip) | (ratio < 1-clip)`. Important: it is NOT `|ratio-1| > clip` , that form disagrees on floating-point boundary values (e.g. `exp(log(1.2))` = 1.20000004 counts as clipped under abs-form but not under the range-form the oracle uses).\n- `bradley_terry_logit = beta * (chosen_reward - rejected_reward)` , a plain logit, no sigmoid/logsigmoid. Elementwise.\n\nVerification: 452 comparisons across 30+ random trials with diverse shapes `(4,), (3,5), (2,3,4), (1,), (10,), (7,2)`, all three estimators, `clip=None` and clip values `{0.0,0.05,0.1,0.2,0.3,0.5,1.0}`, boundary ratios exactly at `1±clip`, and inf/nan inputs , all pass `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)`."}

{"type": "text", "text": "agentId: a5acd56804a9a0d15 (use SendMessage with to: 'a5acd56804a9a0d15', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 29330\ntool_uses: 33\nduration_ms: 212902</usage>"}

{"type": "text", "text": "All 10818 randomized trials pass across diverse shapes (1D/2D/3D), mask kinds (full/zero/sparse/half), all axes, both shift modes, multiple eps values, and empty masks. Combined with the earlier suites and dtype tests, the implementation exactly reproduces `reforacle`.\n\nRecovered conventions:\n- **masked_mean**: `sum(values*mask)/sum(mask)` (mask = float weights), safe-divides to 0 on empty mask; respects `axis` with the reduced dim dropped.\n- **masked_sum**: `sum(values*mask)` over axis (or global).\n- **masked_max**: global scalar, excludes masked positions via `where(mask, values, -inf)`; empty mask returns 0 in the input dtype.\n- **whiten**: masked biased mean/var (÷sum(mask)), `(x-mean)/sqrt(var+1e-8)` applied to all elements; if `shift_mean` is False, adds mean back. No output masking.\n- **masked_whiten**: same as whiten but multiplies the whitened result by mask before optionally adding mean back.\n- **advantage_mean_std**: returns `stack([masked_mean, sqrt(masked_biased_var)])` , biased std, no eps, empty→[0,0].\n- **normalize**: global `(x - x.mean()) / (x.std(unbiased=True) + eps)` , unbiased std, eps added to std (not inside sqrt).\n\nFinal implementation:\n\n```python\nimport torch\n\ndef masked_mean(values, mask, axis=None):\n    if axis is None:\n        s = (values * mask).sum()\n        d = mask.sum()\n    else:\n        s = (values * mask).sum(dim=axis)\n        d = mask.sum(dim=axis)\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef masked_sum(values, mask, axis=None):\n    if axis is None:\n        return (values * mask).sum()\n    return (values * mask).sum(dim=axis)\n\ndef masked_max(values, mask):\n    neg_inf = torch.finfo(values.dtype).min\n    masked_vals = torch.where(mask.bool(), values, torch.full_like(values, neg_inf))\n    if mask.sum() == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return masked_vals.max()\n\ndef _maskprim_safe_mean(values, mask):\n    s = (values * mask).sum()\n    d = mask.sum()\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef _maskprim_mvar(values, mask):\n    mean = _maskprim_safe_mean(values, mask)\n    centered = values - mean\n    var = _maskprim_safe_mean(centered * centered, mask)\n    return mean, var\n\ndef whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef masked_whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    whitened = whitened * mask\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef advantage_mean_std(advantages, mask):\n    mean, var = _maskprim_mvar(advantages, mask)\n    std = torch.sqrt(var)\n    return torch.stack([mean, std])\n\ndef normalize(x, eps):\n    mean = x.mean()\n    std = x.std(unbiased=True)\n    return (x - mean) / (std + eps)\n```"}

{"type": "text", "text": "agentId: af76df31e94c64aa5 (use SendMessage with to: 'af76df31e94c64aa5', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 33906\ntool_uses: 19\nduration_ms: 267714</usage>"}

{"type": "text", "text": "All final edge cases pass, including float64 dtype preservation. The implementation is verified. Here is the final standalone source.\n\n```python\nimport torch\n\n\ndef top_p_mask(probs, p):\n    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)\n    cumsum = torch.cumsum(sorted_probs, dim=-1)\n    cumsum_before = cumsum - sorted_probs\n    keep = cumsum_before < p\n    keep[..., 0] = True\n    kept = sorted_probs * keep.to(sorted_probs.dtype)\n    kept = kept / kept.sum(dim=-1, keepdim=True)\n    out = torch.zeros_like(kept)\n    out.scatter_(-1, sorted_idx, kept)\n    return out\n\n\ndef top_k_mask(logits, k):\n    kth = torch.topk(logits, k, dim=-1).values[..., -1:]\n    return logits >= kth\n\n\ndef argmax_tokens(logits):\n    flipped = torch.flip(logits, dims=[-1])\n    n = logits.shape[-1]\n    idx = torch.argmax(flipped, dim=-1)\n    return (n - 1) - idx\n\n\ndef mode_label(labels):\n    def _mode_1d(v):\n        uniq, counts = torch.unique(v, return_counts=True)\n        max_count = counts.max()\n        candidates = uniq[counts == max_count]\n        return candidates.max()\n\n    if labels.dim() == 1:\n        return _mode_1d(labels)\n    return torch.stack([_mode_1d(labels[i]) for i in range(labels.shape[0])])\n\n\ndef _positional_quantile(x, q, interpolation):\n    xs, _ = torch.sort(x, dim=-1)\n    n = xs.shape[0]\n    pos = (n - 1) * q\n    if interpolation == \"lower\":\n        return xs[int(pos // 1)]\n    lo = int(pos // 1)\n    hi = min(lo + 1, n - 1)\n    frac = pos - lo\n    return xs[lo] * (1 - frac) + xs[hi] * frac\n\n\ndef median_reward(rewards):\n    return _positional_quantile(rewards, 0.5, \"linear\")\n\n\ndef quantile_lower(x, q):\n    return _positional_quantile(x, q, \"lower\")\n\n\ndef pad_mask_from_lengths(lengths, max_len):\n    arange = torch.arange(max_len, device=lengths.device).unsqueeze(0)\n    return (arange < lengths.unsqueeze(1)).to(torch.float32)\n\n\ndef first_nonzero_index(mask):\n    nz = mask != 0\n    any_nz = nz.any(dim=-1)\n    idx = torch.argmax(nz.to(torch.long), dim=-1)\n    return torch.where(any_nz, idx, torch.full_like(idx, -1))\n\n\ndef cumulative_max(x):\n    out = torch.empty_like(x)\n    out[0] = x[0]\n    if x.shape[0] > 1:\n        out[1:] = torch.cummax(x, dim=0)[0][:-1]\n    return out\n\n\ndef bucketize_reward(rewards, boundaries):\n    return torch.bucketize(rewards, boundaries, right=True)\n```\n\n## Summary\n\nI reverse-engineered all 10 functions by probing `reforacle` and verified against 100+ diverse trials (including heavy tie/edge cases) with zero mismatches. Key recovered conventions, several of which are non-obvious:\n\n- **top_p_mask**: nucleus over last dim; keep where `cumsum_before < p` (HF shift, crossing token included), always keeps ≥1 token, output is **renormalized probabilities** in original order (not a boolean mask), zeros elsewhere.\n- **top_k_mask**: **threshold-based** (`logit >= kth-largest value`), so ties at the threshold are all kept and the count may exceed k. Returns bool, last dim.\n- **argmax_tokens**: last dim, but returns the **LAST** maximum index on ties (not torch.argmax's first).\n- **mode_label**: most frequent value per row (dim 0 for 1D, per-row for 2D); tie-break returns the **largest** value (not torch.mode which returns smallest). Handles negatives.\n- **median_reward / quantile_lower**: reduce along dim 0, but with a non-standard rule , **sort along the last dim, then take the positional quantile across dim 0 without re-sorting** (median=linear interp, quantile_lower=floor/lower). Standard `torch.quantile(dim=0)` does NOT reproduce this because it re-sorts each column. For 1D this coincides with normal quantile.\n- **pad_mask_from_lengths**: `(N, max_len)` float32, `arange < length`.\n- **first_nonzero_index**: per-row along last dim, returns **-1** when a row is all zero.\n- **cumulative_max**: exclusive right-shifted running max along **dim 0** (`out[0]=x[0]`, `out[i]=max(x[0..i-1])`).\n- **bucketize_reward**: `torch.bucketize(..., right=True)`.\n\nFiles: `/app/myimpl.py` (final implementation), verification scripts `/app/verify2.py`, `/app/verify4.py`, `/app/verify5.py`. The `/app/objectives.py` scaffold was left unmodified (it contains many other stubs unrelated to this task)."}

{"type": "text", "text": "agentId: aca416a9a99d1536f (use SendMessage with to: 'aca416a9a99d1536f', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 63559\ntool_uses: 37\nduration_ms: 552401</usage>"}

{
  "description": "Reverse-engineer tensor-op primitives",
  "prompt": "You are reverse-engineering a sealed reference module `reforacle` (importable via `import reforacle`, source unreadable). Implement Python functions reproducing `reforacle`'s EXACT numerical behavior. ONLY dependency may be `torch`. Recover behavior ONLY by calling reforacle and comparing , never decompile.\n\nRun `python3 /app/probe.py` to see calling pattern. Oracle called like `reforacle.cumulative_max(x)`.\n\nImplement EXACTLY these functions (identical signatures):\n\ndef top_p_mask(probs, p):\ndef top_k_mask(logits, k):\ndef argmax_tokens(logits):\ndef mode_label(labels):\ndef median_reward(rewards):\ndef quantile_lower(x, q):\ndef pad_mask_from_lengths(lengths, max_len):\ndef first_nonzero_index(mask):\ndef cumulative_max(x):\ndef bucketize_reward(rewards, boundaries):\n\nKNOWN probed facts (verify and build on): \n- cumulative_max([1,3,2,5,4]) -> [1,1,3,3,5]. This is an EXCLUSIVE running max shifted right: out[0]=x[0], out[i]=max(x[0..i-1]) for i>=1. VERIFY this thoroughly on 2D too (which axis? probably last dim or dim 0 , probe 2D input shapes).\n- median_reward([1,2,3,4]) -> 2.5 (mean of two middle elements, numpy-style, NOT torch.median). Verify odd/even, 2D behavior.\n- quantile_lower([1,2,3,4],0.5) -> 2.0 (torch.quantile interpolation='lower'). Verify with various q and shapes.\n- argmax_tokens(randn(2,3,4)) -> shape (2,3): argmax over LAST dim. Verify tie-breaking (torch.argmax returns first max).\n\nProbe extensively for exact conventions:\n- top_p_mask(probs, p): nucleus mask. Returns boolean/float mask over vocab of tokens kept in top-p nucleus. Determine: sort descending, cumulative sum, keep tokens until cumsum>=p. Is the token that crosses p INCLUDED (standard: yes)? Is at least 1 token always kept? Is the mask in original order or sorted order? dtype of output (bool/float)? Probe with crafted probs like [0.5,0.3,0.15,0.05], p=0.6,0.8,0.9. Determine whether comparison is cumsum > p (shift) or >=. Standard HF: keep where cumprobs - probs < p i.e. shift right so first token beyond threshold is kept. NAIL this exactly.\n- top_k_mask(logits, k): mask keeping top-k logits. Output dtype (bool/float)? Original order? Which axis (last)? What is value at kept vs removed (True/1 vs False/0, or logits vs -inf)? Probe. Ties?\n- mode_label(labels): most frequent label. Return scalar? For ties which is returned (torch.mode returns...)? Probe with ties. 2D behavior (per row? which dim?).\n- pad_mask_from_lengths(lengths, max_len): lengths shape (N,), returns (N, max_len) mask where positions < length are 1(True). Determine dtype (bool/float) and whether it's position<length. Probe.\n- first_nonzero_index(mask): index of first nonzero along... which axis? Returns per-row index? What if all zero (returns 0? -1? max_len?)? Probe 1D and 2D. dtype long.\n- bucketize_reward(rewards, boundaries): torch.bucketize. Determine right=True or False (boundary inclusion). Probe values exactly on boundaries. Output dtype long.\n\nVERIFICATION: Write a test comparing impl to reforacle over >=30 trials each with diverse shapes/values and EDGE CASES (ties, all-zero mask, values exactly on boundaries, k=1, k=all, p=0, p=1, single element). For float outputs use torch.allclose(atol=1e-6,rtol=1e-5,equal_nan=True); for int/bool outputs use torch.equal. Iterate until all pass.\n\nReturn ONLY final Python source for these 10 functions (no prose/markdown), standalone torch-only (`import torch` assumed). Helpers nested or `_top_` prefixed. No reforacle reference in returned code.",
  "subagent_type": "general-purpose"
}

, output, 
{"type": "text", "text": "All 60 trials pass across scalar and vector shapes, diverse beta, and label_smoothing (including 0).\n\n```python\ndef dpo_loss(pc, pr, rc, rr, beta, label_smoothing):\n    import torch\n    logits = beta * ((pc - rc) - (pr - rr))\n    loss = (\n        -(1 - label_smoothing) * torch.nn.functional.logsigmoid(logits)\n        - label_smoothing * torch.nn.functional.logsigmoid(-logits)\n    )\n    return loss.mean()\n\n\ndef ipo_loss(pc, pr, rc, rr, beta):\n    h = (pc - rc) - (pr - rr)\n    loss = (h - 1 / (2 * beta)) ** 2\n    return loss.mean()\n```\n\nKey facts for composite objectives:\n\n- DPO logits formula: `logits = beta * ((pc - rc) - (pr - rr))`\n- DPO label_smoothing formula (per-example, before reduction): `-(1 - label_smoothing) * logsigmoid(logits) - label_smoothing * logsigmoid(-logits)`. At `label_smoothing=0` this collapses to `-logsigmoid(logits)`.\n- IPO formula (per-example): `((pc - rc) - (pr - rr) - 1/(2*beta))**2`, i.e. `(h - 1/(2*beta))**2` with `h = (pc-rc)-(pr-rr)`. Beta appears only inside the `1/(2*beta)` margin term, not as a multiplier on `h`.\n- Reduction convention: BOTH functions return a MEAN scalar (shape `[]`), not per-example. Verified by passing N=4 vectors and getting a scalar equal to the per-example `.mean()`.\n\nVerified over 60 trials (scalar and `(N,)` shapes, beta in ~[0.05, 2.05], label_smoothing including 0) with `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)` , all pass."}

{"type": "text", "text": "agentId: ab14fc7203cb72eaf (use SendMessage with to: 'ab14fc7203cb72eaf', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 12305\ntool_uses: 3\nduration_ms: 37779</usage>"}

{"type": "text", "text": "All 60 trials plus edge cases pass with zero failures. Here is the final implementation.\n\n```python\nimport torch\n\n\ndef _lp_masked_mean(values, mask):\n    denom = mask.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return (values * mask).sum() / denom\n\n\ndef logsumexp(x, axis):\n    return torch.logsumexp(x, dim=axis)\n\n\ndef log_softmax(x, axis):\n    return x - torch.logsumexp(x, dim=axis, keepdim=True)\n\n\ndef token_logprobs(logits, labels):\n    lsm = torch.log_softmax(logits, dim=-1)\n    return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n\n\ndef selective_logprobs(logits, labels, mask):\n    lp = token_logprobs(logits, labels)\n    return (lp * mask).sum(dim=-1)\n\n\ndef sequence_logprob(logits, labels, mask, length_normalize):\n    lp = token_logprobs(logits, labels)\n    total = (lp * mask).sum(dim=-1)\n    if length_normalize:\n        denom = mask.sum(dim=-1)\n        safe = torch.where(denom > 0, denom, torch.ones_like(denom))\n        out = total / safe\n        return torch.where(denom > 0, out, torch.zeros_like(out))\n    return total\n\n\ndef logprob_at_temperature(logits, labels, temperature):\n    return token_logprobs(logits / temperature, labels)\n\n\ndef cross_entropy(logits, labels, ignore_index):\n    labels = labels.long()\n    flat_logits = logits.reshape(-1, logits.shape[-1])\n    flat_labels = labels.reshape(-1)\n    lsm = torch.log_softmax(flat_logits, dim=-1)\n    valid = (flat_labels != ignore_index)\n    safe_labels = torch.where(valid, flat_labels, torch.zeros_like(flat_labels))\n    nll = -lsm.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)\n    valid_f = valid.to(nll.dtype)\n    denom = valid_f.sum()\n    if denom == 0:\n        return torch.zeros((), dtype=nll.dtype, device=nll.device)\n    return (nll * valid_f).sum() / denom\n\n\ndef smoothed_nll(logits, labels, smoothing):\n    lsm = torch.log_softmax(logits, dim=-1)\n    nll = -lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n    smooth = -lsm.mean(dim=-1)\n    loss = (1.0 - smoothing) * nll + smoothing * smooth\n    return loss.mean()\n\n\ndef entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    return _lp_masked_mean(ent, mask)\n\n\ndef normalized_entropy(logits, mask):\n    lsm = torch.log_softmax(logits, dim=-1)\n    p = torch.softmax(logits, dim=-1)\n    ent = -(p * lsm).sum(dim=-1)\n    vocab = logits.shape[-1]\n    norm = ent / torch.log(torch.tensor(float(vocab), dtype=ent.dtype, device=ent.device))\n    return _lp_masked_mean(norm, mask)\n```\n\nRecovered conventions (all verified over 60 diverse trials + edge cases, 0 failures at atol=1e-6/rtol=1e-5):\n- **logsumexp**: standard, no keepdim; **log_softmax**: `x - logsumexp(x, keepdim=True)`.\n- **token_logprobs**: log_softmax over last dim, gather labels; output shape = labels shape.\n- **selective_logprobs**: `sum(token_logprobs * mask, dim=-1)` , sums over the sequence (last) axis, dropping it.\n- **sequence_logprob**: same masked sum; `length_normalize=True` divides by `mask.sum(-1)`, returning 0 when the mask is all-zero (guarded div).\n- **logprob_at_temperature**: divides logits by temperature before log_softmax.\n- **cross_entropy**: flattens to (-1, V), mean of NLL over non-ignored positions (`sum/count`); returns 0 (not NaN) when all positions are ignored. Works for any ignore_index value.\n- **smoothed_nll**: `mean[(1-s)*nll + s*(-mean_vocab log_softmax)]`.\n- **entropy**: per-position natural-log entropy `-sum(p*log p)`, masked-mean over positions.\n- **normalized_entropy**: entropy / log(vocab_size), masked-mean; returns 0 on empty mask.\n\nNote: `masked_mean` on an all-zero mask returns 0, which is the guard reused by entropy/normalized_entropy."}

{"type": "text", "text": "agentId: acac49472241baa47 (use SendMessage with to: 'acac49472241baa47', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 26571\ntool_uses: 19\nduration_ms: 171626</usage>"}

{"type": "text", "text": "All 1762 tests pass across diverse shapes, group sizes, gamma/lam values, scale_by_std both ways, clip values, and boundary cases.\n\n```python\nimport torch\n\n\ndef _rl_masked_mean(x, mask):\n    return (x * mask).sum() / mask.sum()\n\n\ndef grpo_advantages(rewards, group_size, scale_by_std):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    if scale_by_std:\n        std = groups.std(dim=1, unbiased=False, keepdim=True)\n        adv = adv / (std + 1e-6)\n    return adv.reshape(-1)\n\n\ndef rloo_advantages(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    group_sum = groups.sum(dim=1, keepdim=True)\n    baseline = (group_sum - groups) / (group_size - 1)\n    adv = groups - baseline\n    return adv.reshape(-1)\n\n\ndef group_mean_baseline(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    mean = groups.mean(dim=1, keepdim=True)\n    adv = groups - mean\n    return adv.reshape(-1)\n\n\ndef gae(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    adv = torch.zeros_like(rewards)\n    gae_acc = torch.zeros_like(rewards[0])\n    next_v = next_value\n    for t in range(T - 1, -1, -1):\n        delta = rewards[t] + gamma * next_v - values[t]\n        gae_acc = delta + gamma * lam * gae_acc\n        adv[t] = gae_acc\n        next_v = values[t]\n    return adv\n\n\ndef discounted_returns(rewards, gamma):\n    T = rewards.shape[0]\n    out = torch.zeros_like(rewards)\n    acc = torch.zeros_like(rewards[0])\n    for t in range(T - 1, -1, -1):\n        acc = rewards[t] + gamma * acc\n        out[t] = acc\n    return out\n\n\ndef lambda_returns(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    out = torch.zeros_like(rewards)\n    next_ret = next_value\n    for t in range(T - 1, -1, -1):\n        next_v = next_value if t == T - 1 else values[t + 1]\n        next_ret = rewards[t] + gamma * ((1 - lam) * next_v + lam * next_ret)\n        out[t] = next_ret\n    return out\n\n\ndef clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):\n    ratio = torch.exp(logp - old_logp)\n    unclipped = ratio * advantages\n    clipped = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) * advantages\n    per = -torch.min(unclipped, clipped)\n    return _rl_masked_mean(per, mask)\n\n\ndef value_loss(values, old_values, returns, clip):\n    unclipped = (values - returns) ** 2\n    vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)\n    clipped = (vpred_clipped - returns) ** 2\n    return 0.5 * torch.max(unclipped, clipped).mean()\n\n\ndef huber_value_loss(values, returns, delta):\n    x = values - returns\n    absx = x.abs()\n    huber = torch.where(absx <= delta, 0.5 * x ** 2, delta * (absx - 0.5 * delta))\n    return huber.mean()\n```\n\nKey conventions you asked about for composite objectives:\n\n- **GAE / returns time-axis & shape**: Time is the **first axis (dim 0)**. Inputs may be 1-D `(T,)` or N-D `(T, ...)` (e.g. `(T, B)`); the recursion runs over dim 0 and per-timestep slices broadcast over trailing dims. `next_value` may be a scalar or shape `(...)` matching the trailing dims. `gae` returns **advantages only** (a single tensor, not a tuple), same shape as `rewards`. `lambda_returns` returns TD(λ) returns and exactly equals `gae + values`. `discounted_returns` also reduces along dim 0.\n\n- **grpo_advantages / group_mean_baseline shape**: Input is flattened then reshaped to `(-1, group_size)` (row-major grouping), and **output is always flat `(N,)`** regardless of input shape. std is **biased** (`unbiased=False`), eps = **1e-6**. Note `group_mean_baseline` returns the **advantages** (`reward - group_mean`), not the raw baseline.\n\n- **clipped_pg_loss sign & clamp**: `ratio = exp(logp - old_logp)`, `clamp(ratio, 1 - clip_low, 1 + clip_high)` (clip_low is the lower bound, clip_high the upper), then **`-min(unclipped, clipped)`** (negative , it is a loss), reduced by **masked mean = `sum(per*mask)/sum(mask)`**.\n\n- **value_loss 0.5 factor**: Yes , `0.5 * mean(max(unclipped, clipped))`, using **max** of clipped vs unclipped squared errors, plain `.mean()` over all elements. `clip` is always numeric (the oracle does not special-case `None`; passing `None` errors). `clip=0.0` degenerates to `0.5*mean((old_values - returns)^2)` vs unclipped max.\n\n- **huber_value_loss**: standard smooth-L1 (`0.5*x^2` for `|x|<=delta`, else `delta*(|x|-0.5*delta)`), reduced by plain mean; equals `F.smooth_l1_loss(..., beta=delta)`."}

{"type": "text", "text": "agentId: ae12d609b60dce0b3 (use SendMessage with to: 'ae12d609b60dce0b3', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 25299\ntool_uses: 11\nduration_ms: 165387</usage>"}

{"type": "text", "text": "All edge cases pass. Implementation complete and verified in `/app/objectives.py`.\n\nFinal Python source for the 6 functions:\n\n```python\nimport torch\n\n\ndef kl_penalty(logp, ref_logp, estimator):\n    logratio = logp - ref_logp\n    if estimator == 'k1':\n        return logratio\n    elif estimator == 'k2':\n        return 0.5 * logratio * logratio\n    elif estimator == 'k3':\n        return torch.exp(-logratio) + logratio - 1\n    else:\n        raise ValueError(estimator)\n\n\ndef reverse_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    return torch.exp(-logratio) + logratio - 1\n\n\ndef symmetric_kl(logp, ref_logp):\n    logratio = logp - ref_logp\n    reverse = torch.exp(-logratio) + logratio - 1\n    forward = torch.exp(logratio) - logratio - 1\n    return 0.5 * (reverse + forward)\n\n\ndef importance_ratio(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    if clip is not None:\n        ratio = torch.clamp(ratio, 1 - clip, 1 + clip)\n    return ratio\n\n\ndef clip_fraction(logp, old_logp, clip):\n    ratio = torch.exp(logp - old_logp)\n    return ((ratio > 1 + clip) | (ratio < 1 - clip)).float().mean()\n\n\ndef bradley_terry_logit(chosen_reward, rejected_reward, beta):\n    return beta * (chosen_reward - rejected_reward)\n```\n\nNote on conventions (for composite objectives):\n\n- Supported estimator strings for `kl_penalty`: exactly `'k1'`, `'k2'`, `'k3'`. Any other string raises `ValueError(estimator)`.\n- Sign convention: `logratio = logp - ref_logp` (policy minus reference). Then:\n  - `k1 = logratio` (elementwise, no reduction, same shape as input)\n  - `k2 = 0.5 * logratio**2`\n  - `k3 = exp(-logratio) + logratio - 1` (the non-negative unbiased estimator, using `-logratio` inside exp)\n- `reverse_kl` equals the k3 form: `exp(-logratio) + logratio - 1`. Elementwise, no reduction.\n- `symmetric_kl = 0.5 * (reverse + forward)` where `reverse = exp(-logratio)+logratio-1` and `forward = exp(logratio)-logratio-1`. Elementwise.\n- `importance_ratio`: `ratio = exp(logp - old_logp)`; `clip=None` means no clamp; otherwise clamp to `[1-clip, 1+clip]`. Elementwise.\n- `clip_fraction`: scalar (0-dim tensor) mean of the strict out-of-range condition `(ratio > 1+clip) | (ratio < 1-clip)`. Important: it is NOT `|ratio-1| > clip` , that form disagrees on floating-point boundary values (e.g. `exp(log(1.2))` = 1.20000004 counts as clipped under abs-form but not under the range-form the oracle uses).\n- `bradley_terry_logit = beta * (chosen_reward - rejected_reward)` , a plain logit, no sigmoid/logsigmoid. Elementwise.\n\nVerification: 452 comparisons across 30+ random trials with diverse shapes `(4,), (3,5), (2,3,4), (1,), (10,), (7,2)`, all three estimators, `clip=None` and clip values `{0.0,0.05,0.1,0.2,0.3,0.5,1.0}`, boundary ratios exactly at `1±clip`, and inf/nan inputs , all pass `torch.allclose(atol=1e-6, rtol=1e-5, equal_nan=True)`."}

{"type": "text", "text": "agentId: a5acd56804a9a0d15 (use SendMessage with to: 'a5acd56804a9a0d15', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 29330\ntool_uses: 33\nduration_ms: 212902</usage>"}

{"type": "text", "text": "All 10818 randomized trials pass across diverse shapes (1D/2D/3D), mask kinds (full/zero/sparse/half), all axes, both shift modes, multiple eps values, and empty masks. Combined with the earlier suites and dtype tests, the implementation exactly reproduces `reforacle`.\n\nRecovered conventions:\n- **masked_mean**: `sum(values*mask)/sum(mask)` (mask = float weights), safe-divides to 0 on empty mask; respects `axis` with the reduced dim dropped.\n- **masked_sum**: `sum(values*mask)` over axis (or global).\n- **masked_max**: global scalar, excludes masked positions via `where(mask, values, -inf)`; empty mask returns 0 in the input dtype.\n- **whiten**: masked biased mean/var (÷sum(mask)), `(x-mean)/sqrt(var+1e-8)` applied to all elements; if `shift_mean` is False, adds mean back. No output masking.\n- **masked_whiten**: same as whiten but multiplies the whitened result by mask before optionally adding mean back.\n- **advantage_mean_std**: returns `stack([masked_mean, sqrt(masked_biased_var)])` , biased std, no eps, empty→[0,0].\n- **normalize**: global `(x - x.mean()) / (x.std(unbiased=True) + eps)` , unbiased std, eps added to std (not inside sqrt).\n\nFinal implementation:\n\n```python\nimport torch\n\ndef masked_mean(values, mask, axis=None):\n    if axis is None:\n        s = (values * mask).sum()\n        d = mask.sum()\n    else:\n        s = (values * mask).sum(dim=axis)\n        d = mask.sum(dim=axis)\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef masked_sum(values, mask, axis=None):\n    if axis is None:\n        return (values * mask).sum()\n    return (values * mask).sum(dim=axis)\n\ndef masked_max(values, mask):\n    neg_inf = torch.finfo(values.dtype).min\n    masked_vals = torch.where(mask.bool(), values, torch.full_like(values, neg_inf))\n    if mask.sum() == 0:\n        return torch.zeros((), dtype=values.dtype, device=values.device)\n    return masked_vals.max()\n\ndef _maskprim_safe_mean(values, mask):\n    s = (values * mask).sum()\n    d = mask.sum()\n    safe = d + (d == 0).to(d.dtype)\n    return s / safe * (d != 0).to(s.dtype)\n\ndef _maskprim_mvar(values, mask):\n    mean = _maskprim_safe_mean(values, mask)\n    centered = values - mean\n    var = _maskprim_safe_mean(centered * centered, mask)\n    return mean, var\n\ndef whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef masked_whiten(values, mask, shift_mean):\n    mean, var = _maskprim_mvar(values, mask)\n    whitened = (values - mean) / torch.sqrt(var + 1e-8)\n    whitened = whitened * mask\n    if not shift_mean:\n        whitened = whitened + mean\n    return whitened\n\ndef advantage_mean_std(advantages, mask):\n    mean, var = _maskprim_mvar(advantages, mask)\n    std = torch.sqrt(var)\n    return torch.stack([mean, std])\n\ndef normalize(x, eps):\n    mean = x.mean()\n    std = x.std(unbiased=True)\n    return (x - mean) / (std + eps)\n```"}

{"type": "text", "text": "agentId: af76df31e94c64aa5 (use SendMessage with to: 'af76df31e94c64aa5', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 33906\ntool_uses: 19\nduration_ms: 267714</usage>"}

{"type": "text", "text": "All final edge cases pass, including float64 dtype preservation. The implementation is verified. Here is the final standalone source.\n\n```python\nimport torch\n\n\ndef top_p_mask(probs, p):\n    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)\n    cumsum = torch.cumsum(sorted_probs, dim=-1)\n    cumsum_before = cumsum - sorted_probs\n    keep = cumsum_before < p\n    keep[..., 0] = True\n    kept = sorted_probs * keep.to(sorted_probs.dtype)\n    kept = kept / kept.sum(dim=-1, keepdim=True)\n    out = torch.zeros_like(kept)\n    out.scatter_(-1, sorted_idx, kept)\n    return out\n\n\ndef top_k_mask(logits, k):\n    kth = torch.topk(logits, k, dim=-1).values[..., -1:]\n    return logits >= kth\n\n\ndef argmax_tokens(logits):\n    flipped = torch.flip(logits, dims=[-1])\n    n = logits.shape[-1]\n    idx = torch.argmax(flipped, dim=-1)\n    return (n - 1) - idx\n\n\ndef mode_label(labels):\n    def _mode_1d(v):\n        uniq, counts = torch.unique(v, return_counts=True)\n        max_count = counts.max()\n        candidates = uniq[counts == max_count]\n        return candidates.max()\n\n    if labels.dim() == 1:\n        return _mode_1d(labels)\n    return torch.stack([_mode_1d(labels[i]) for i in range(labels.shape[0])])\n\n\ndef _positional_quantile(x, q, interpolation):\n    xs, _ = torch.sort(x, dim=-1)\n    n = xs.shape[0]\n    pos = (n - 1) * q\n    if interpolation == \"lower\":\n        return xs[int(pos // 1)]\n    lo = int(pos // 1)\n    hi = min(lo + 1, n - 1)\n    frac = pos - lo\n    return xs[lo] * (1 - frac) + xs[hi] * frac\n\n\ndef median_reward(rewards):\n    return _positional_quantile(rewards, 0.5, \"linear\")\n\n\ndef quantile_lower(x, q):\n    return _positional_quantile(x, q, \"lower\")\n\n\ndef pad_mask_from_lengths(lengths, max_len):\n    arange = torch.arange(max_len, device=lengths.device).unsqueeze(0)\n    return (arange < lengths.unsqueeze(1)).to(torch.float32)\n\n\ndef first_nonzero_index(mask):\n    nz = mask != 0\n    any_nz = nz.any(dim=-1)\n    idx = torch.argmax(nz.to(torch.long), dim=-1)\n    return torch.where(any_nz, idx, torch.full_like(idx, -1))\n\n\ndef cumulative_max(x):\n    out = torch.empty_like(x)\n    out[0] = x[0]\n    if x.shape[0] > 1:\n        out[1:] = torch.cummax(x, dim=0)[0][:-1]\n    return out\n\n\ndef bucketize_reward(rewards, boundaries):\n    return torch.bucketize(rewards, boundaries, right=True)\n```\n\n## Summary\n\nI reverse-engineered all 10 functions by probing `reforacle` and verified against 100+ diverse trials (including heavy tie/edge cases) with zero mismatches. Key recovered conventions, several of which are non-obvious:\n\n- **top_p_mask**: nucleus over last dim; keep where `cumsum_before < p` (HF shift, crossing token included), always keeps ≥1 token, output is **renormalized probabilities** in original order (not a boolean mask), zeros elsewhere.\n- **top_k_mask**: **threshold-based** (`logit >= kth-largest value`), so ties at the threshold are all kept and the count may exceed k. Returns bool, last dim.\n- **argmax_tokens**: last dim, but returns the **LAST** maximum index on ties (not torch.argmax's first).\n- **mode_label**: most frequent value per row (dim 0 for 1D, per-row for 2D); tie-break returns the **largest** value (not torch.mode which returns smallest). Handles negatives.\n- **median_reward / quantile_lower**: reduce along dim 0, but with a non-standard rule , **sort along the last dim, then take the positional quantile across dim 0 without re-sorting** (median=linear interp, quantile_lower=floor/lower). Standard `torch.quantile(dim=0)` does NOT reproduce this because it re-sorts each column. For 1D this coincides with normal quantile.\n- **pad_mask_from_lengths**: `(N, max_len)` float32, `arange < length`.\n- **first_nonzero_index**: per-row along last dim, returns **-1** when a row is all zero.\n- **cumulative_max**: exclusive right-shifted running max along **dim 0** (`out[0]=x[0]`, `out[i]=max(x[0..i-1])`).\n- **bucketize_reward**: `torch.bucketize(..., right=True)`.\n\nFiles: `/app/myimpl.py` (final implementation), verification scripts `/app/verify2.py`, `/app/verify4.py`, `/app/verify5.py`. The `/app/objectives.py` scaffold was left unmodified (it contains many other stubs unrelated to this task)."}

{"type": "text", "text": "agentId: aca416a9a99d1536f (use SendMessage with to: 'aca416a9a99d1536f', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 63559\ntool_uses: 37\nduration_ms: 552401</usage>"}

{
  "description": "Reverse-engineer composite objectives",
  "prompt": "You are reverse-engineering a sealed reference module `reforacle` (importable via `import reforacle`, source unreadable). Implement Python functions reproducing `reforacle`'s EXACT numerical behavior. ONLY dependency may be `torch`. Recover behavior ONLY by calling reforacle and comparing , never decompile. Tolerance for these composite objectives is atol=1e-5, rtol=1e-4.\n\nRun `python3 /app/probe.py` for calling pattern. Oracle called like `reforacle.grpo_objective(...)`.\n\nImplement EXACTLY these 5 functions (identical signatures):\n\ndef dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,\n                      chosen_labels, rejected_labels, chosen_mask, rejected_mask,\n                      beta, label_smoothing):\n\ndef grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,\n                   rewards, group_size, beta, clip_low, clip_high, scale_by_std,\n                   kl_estimator):\n\ndef ppo_objective(rewards, values, old_values, logp, old_logp, next_value,\n                  gamma, lam, clip_low, clip_high, vf_clip, vf_coef):\n\ndef rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):\n\ndef reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):\n\nThese compose lower-level quantities. I have ALREADY recovered the exact conventions of reforacle's primitives (which these objectives are built from). USE THESE EXACT CONVENTIONS , they are ground truth, but you must still PROBE the composite functions to determine how the pieces combine (signs, coefficients, reductions, KL term placement):\n\nRECOVERED PRIMITIVE CONVENTIONS (reforacle uses these internally):\n- token_logprobs(logits, labels): lsm = log_softmax(logits, dim=-1); return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1). Output shape = labels shape. Softmax over LAST dim.\n- sequence_logprob(logits, labels, mask, length_normalize): per-token logprobs summed over last dim with mask: (lp*mask).sum(-1); if length_normalize: divide by mask.sum(-1) (guard 0).\n- selective_logprobs = (token_logprobs * mask).sum(dim=-1).\n- masked_mean(values, mask) = (values*mask).sum()/mask.sum() (global; 0 if empty).\n- clipped_pg_loss: ratio=exp(logp-old_logp); unclipped=ratio*adv; clipped=clamp(ratio, 1-clip_low, 1+clip_high)*adv; per = -min(unclipped, clipped); reduce = masked_mean over mask = sum(per*mask)/sum(mask).\n- grpo_advantages(rewards, group_size, scale_by_std): flat=rewards.reshape(-1); groups=flat.reshape(-1, group_size); adv=groups-groups.mean(1,keepdim); if scale_by_std: adv/=(groups.std(1,unbiased=False,keepdim)+1e-6); return adv.reshape(-1). Output flat (N,).\n- rloo_advantages(rewards, group_size): groups=flat.reshape(-1,group_size); baseline=(group_sum-groups)/(group_size-1); adv=groups-baseline; flat output.\n- kl_penalty(logp, ref_logp, estimator): logratio=logp-ref_logp; 'k1'->logratio; 'k2'->0.5*logratio^2; 'k3'->exp(-logratio)+logratio-1. Elementwise.\n- dpo_loss(pc,pr,rc,rr,beta,label_smoothing): logits=beta*((pc-rc)-(pr-rr)); loss = -(1-ls)*logsigmoid(logits) - ls*logsigmoid(-logits); returns loss.mean().\n- gae(rewards,values,next_value,gamma,lam): time axis dim0; delta_t=r_t+gamma*V_{t+1}-V_t (V_T=next_value); A_t=delta_t+gamma*lam*A_{t+1}; returns advantages (dim0). lambda_returns = gae+values.\n- value_loss(values,old_values,returns,clip): unclipped=(values-returns)^2; vpred_clipped=old_values+clamp(values-old_values,-clip,clip); clipped=(vpred_clipped-returns)^2; return 0.5*max(unclipped,clipped).mean().\n- reverse_kl(logp,ref_logp)= exp(-(logp-ref_logp))+(logp-ref_logp)-1 (k3 form). entropy etc as standard.\n\nNOW PROBE the composite functions to determine how they combine. Specifically determine by probing:\n\n1. dpo_sequence_loss: Compute chosen policy seq logprob pc = sequence_logprob(pc_logits, chosen_labels, chosen_mask, length_normalize=?), similarly pr (rejected), rc (ref chosen), rr (ref rejected). Then dpo_loss(pc,pr,rc,rr,beta,label_smoothing). PROBE whether length_normalize is True or False in the seq logprobs (craft masks of differing lengths and compare). PROBE whether logprobs are summed (not length-normalized) , most DPO uses SUM. Determine the mask handling and that it matches dpo_loss reduction (mean).\n\n2. grpo_objective: policy logp = token_logprobs(logits,labels), old_logp = token_logprobs(old_logits,labels), ref_logp = token_logprobs(ref_logits,labels). advantages = grpo_advantages(rewards, group_size, scale_by_std). NOTE rewards is per-sequence (shape (N,)) but logp is per-token (N,T) , determine how advantages broadcast over tokens (each sequence's scalar advantage broadcast to its tokens). pg_loss = clipped_pg_loss(logp, old_logp, advantages_broadcast, completion_mask, clip_low, clip_high). kl = kl_penalty(logp, ref_logp, kl_estimator). total = pg_loss + beta * masked_mean(kl, completion_mask). PROBE: sign of kl term (+beta*kl added to loss), whether kl is masked_mean reduced, whether advantages are broadcast then whitened, exact combination. Determine the return (scalar loss).\n\n3. ppo_objective: advantages = gae(rewards,values,next_value,gamma,lam); returns = advantages+values (lambda_returns). PROBE whether advantages are whitened/normalized before pg loss (common: whiten advantages). pg_loss = clipped_pg_loss(logp, old_logp, advantages, mask=?, clip_low, clip_high) , determine mask (maybe all-ones / mean over all). vf = value_loss(values, old_values, returns, vf_clip). total = pg_loss + vf_coef * vf. PROBE: is there a mask? (signature has no mask -> reduction is plain mean over all elements). Determine whether advantages whitened (try both, compare). Determine exact returns definition.\n\n4. rloo_objective: logp=token_logprobs(logits,labels); old_logp=token_logprobs(old_logits,labels); advantages=rloo_advantages(rewards,group_size) broadcast over tokens; pg_loss=clipped_pg_loss(logp,old_logp,adv,mask,clip_low,clip_high). No KL term (not in signature). Return scalar. PROBE broadcast & reduction.\n\n5. reverse_kl_objective: logp=token_logprobs(logits,labels); ref_logp=token_logprobs(ref_logits,labels). This is a policy-gradient-with-reverse-KL objective. PROBE: likely loss = masked_mean(-advantages*logp, mask) + beta*masked_mean(reverse_kl(logp,ref_logp), mask), OR = -masked_mean(advantages*logp - beta*reverse_kl(...), mask). advantages here is per-token (same shape as logp) , confirm shape. Determine exact sign and combination by crafting inputs where you can compute each candidate and compare.\n\nCRITICAL PROBING METHOD: For each composite, build YOUR candidate implementation using the primitive conventions above, then compare to reforacle over many random trials. When mismatch, isolate the piece (e.g. compare reforacle.grpo_objective vs your value with kl term removed to see the kl coefficient/sign; set beta=0 to isolate pg_loss; set advantages such that things simplify). Use beta=0, clip huge (no clipping), single group, length-1 sequences etc. to isolate each term.\n\nFor broadcasting rewards->tokens in grpo/rloo: rewards shape (N,), logits shape (N,T,V), labels (N,T), completion_mask (N,T). advantages per-sequence (N,) -> unsqueeze(-1) to (N,1) broadcast over T. VERIFY this shape assumption by probing actual shapes reforacle accepts (print shapes in probe).\n\nVERIFICATION: Write a test comparing your impl to reforacle over >=40 random trials each with diverse but VALID shapes, group sizes dividing N, diverse beta (including 0), clip values, both scale_by_std, all kl_estimator strings ('k1','k2','k3'), label_smoothing values. Use torch.allclose(atol=1e-5, rtol=1e-4, equal_nan=True). Iterate until ALL pass including edge cases. First run probe.py-style prints to discover exact input shapes each function expects.\n\nReturn ONLY final Python source for these 5 functions (no prose/markdown fences in the code block, well, provide the code clearly), standalone torch-only (`import torch` assumed at module level). You MAY call the other objective/primitive functions by name IF you also assume they're defined in the same module , but SAFER: make each composite fully self-contained (inline the needed primitive logic) to avoid depending on exact helper availability. Helpers nested or `_obj_` prefixed. No reforacle reference in returned code.",
  "subagent_type": "general-purpose"
}

, output, 
{"type": "text", "text": "The implementation is complete and written to `/app/objectives.py`. All 5 functions verified against `reforacle` over 300+ diverse random trials plus edge cases (empty masks, group_size=1/N, T=1, beta=0, all label_smoothing/kl_estimator/scale_by_std variants) with `atol=1e-5, rtol=1e-4`.\n\nKey findings from probing:\n- **dpo_sequence_loss**: uses `length_normalize=False` (SUM of token logprobs, not mean) for all four sequence logprobs, then standard `dpo_loss` with `.mean()` reduction.\n- **grpo_objective**: standard , per-token clipped PG with per-sequence GRPO advantages broadcast via `unsqueeze(-1)`, plus `beta * masked_mean(kl_penalty)`. KL added (positive sign).\n- **ppo_objective**: GAE advantages are **whitened** (biased std, `unbiased=False`, eps `1e-8`) before the PG loss; `returns = gae + values`; no mask (plain mean over all timesteps); total = `pg + vf_coef * value_loss`.\n- **rloo_objective**: standard , RLOO advantages broadcast over tokens, clipped PG, no KL.\n- **reverse_kl_objective** (the tricky one): NOT a simple `masked_mean(-adv*logp)`. Its two terms are:\n  - PG term = `-(advantages * ((logp*mask).sum(dim=0) / mask.sum())).sum()` , the token logprobs are pooled over the batch dimension (summed over dim 0, divided by the global mask sum) before multiplying by the per-token advantages.\n  - KL term = `N * masked_mean(reverse_kl(logp, ref_logp), mask)` where `N = logits.shape[0]` (a leading batch-size factor), added with coefficient `beta`.\n\nThe 5 required functions are self-contained (all primitives inlined with `_obj_` prefix, torch-only) in `/app/objectives.py`.\n\nFinal source:\n\n```python\nimport torch\n\n\ndef _obj_token_logprobs(logits, labels):\n    lsm = torch.log_softmax(logits, dim=-1)\n    return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)\n\n\ndef _obj_sequence_logprob(logits, labels, mask, length_normalize=False):\n    lp = _obj_token_logprobs(logits, labels)\n    s = (lp * mask).sum(-1)\n    if length_normalize:\n        s = s / mask.sum(-1).clamp(min=1.0)\n    return s\n\n\ndef _obj_masked_mean(values, mask):\n    s = mask.sum()\n    if s > 0:\n        return (values * mask).sum() / s\n    return torch.zeros((), dtype=values.dtype, device=values.device)\n\n\ndef _obj_clipped_pg_loss(logp, old_logp, adv, mask, clip_low, clip_high):\n    ratio = torch.exp(logp - old_logp)\n    unclipped = ratio * adv\n    clipped = torch.clamp(ratio, 1.0 - clip_low, 1.0 + clip_high) * adv\n    per = -torch.min(unclipped, clipped)\n    return _obj_masked_mean(per, mask)\n\n\ndef _obj_grpo_advantages(rewards, group_size, scale_by_std):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    adv = groups - groups.mean(1, keepdim=True)\n    if scale_by_std:\n        adv = adv / (groups.std(1, unbiased=False, keepdim=True) + 1e-6)\n    return adv.reshape(-1)\n\n\ndef _obj_rloo_advantages(rewards, group_size):\n    flat = rewards.reshape(-1)\n    groups = flat.reshape(-1, group_size)\n    gsum = groups.sum(1, keepdim=True)\n    baseline = (gsum - groups) / (group_size - 1)\n    return (groups - baseline).reshape(-1)\n\n\ndef _obj_kl_penalty(logp, ref_logp, estimator):\n    logratio = logp - ref_logp\n    if estimator == \"k1\":\n        return logratio\n    if estimator == \"k2\":\n        return 0.5 * logratio ** 2\n    return torch.exp(-logratio) + logratio - 1.0\n\n\ndef _obj_reverse_kl(logp, ref_logp):\n    lr = logp - ref_logp\n    return torch.exp(-lr) + lr - 1.0\n\n\ndef _obj_dpo_loss(pc, pr, rc, rr, beta, label_smoothing):\n    logits = beta * ((pc - rc) - (pr - rr))\n    loss = (-(1.0 - label_smoothing) * torch.nn.functional.logsigmoid(logits)\n            - label_smoothing * torch.nn.functional.logsigmoid(-logits))\n    return loss.mean()\n\n\ndef _obj_gae(rewards, values, next_value, gamma, lam):\n    T = rewards.shape[0]\n    if not torch.is_tensor(next_value):\n        next_value = torch.tensor(next_value, dtype=values.dtype)\n    adv = torch.zeros_like(rewards)\n    last = torch.zeros_like(rewards[0])\n    next_v = next_value\n    for t in range(T - 1, -1, -1):\n        delta = rewards[t] + gamma * next_v - values[t]\n        last = delta + gamma * lam * last\n        adv[t] = last\n        next_v = values[t]\n    return adv\n\n\ndef _obj_whiten(x, eps=1e-8):\n    return (x - x.mean()) / (x.std(unbiased=False) + eps)\n\n\ndef _obj_value_loss(values, old_values, returns, clip):\n    unclipped = (values - returns) ** 2\n    vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)\n    clipped = (vpred_clipped - returns) ** 2\n    return 0.5 * torch.max(unclipped, clipped).mean()\n\n\ndef dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,\n                      chosen_labels, rejected_labels, chosen_mask, rejected_mask,\n                      beta, label_smoothing):\n    pc = _obj_sequence_logprob(pc_logits, chosen_labels, chosen_mask, length_normalize=False)\n    pr = _obj_sequence_logprob(pr_logits, rejected_labels, rejected_mask, length_normalize=False)\n    rc = _obj_sequence_logprob(rc_logits, chosen_labels, chosen_mask, length_normalize=False)\n    rr = _obj_sequence_logprob(rr_logits, rejected_labels, rejected_mask, length_normalize=False)\n    return _obj_dpo_loss(pc, pr, rc, rr, beta, label_smoothing)\n\n\ndef grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,\n                   rewards, group_size, beta, clip_low, clip_high, scale_by_std,\n                   kl_estimator):\n    logp = _obj_token_logprobs(logits, labels)\n    old_logp = _obj_token_logprobs(old_logits, labels)\n    ref_logp = _obj_token_logprobs(ref_logits, labels)\n    adv = _obj_grpo_advantages(rewards, group_size, scale_by_std).unsqueeze(-1)\n    pg_loss = _obj_clipped_pg_loss(logp, old_logp, adv, completion_mask, clip_low, clip_high)\n    kl = _obj_kl_penalty(logp, ref_logp, kl_estimator)\n    return pg_loss + beta * _obj_masked_mean(kl, completion_mask)\n\n\ndef ppo_objective(rewards, values, old_values, logp, old_logp, next_value,\n                  gamma, lam, clip_low, clip_high, vf_clip, vf_coef):\n    advantages = _obj_gae(rewards, values, next_value, gamma, lam)\n    returns = advantages + values\n    advantages = _obj_whiten(advantages)\n    mask = torch.ones_like(logp)\n    pg_loss = _obj_clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high)\n    vf = _obj_value_loss(values, old_values, returns, vf_clip)\n    return pg_loss + vf_coef * vf\n\n\ndef rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):\n    logp = _obj_token_logprobs(logits, labels)\n    old_logp = _obj_token_logprobs(old_logits, labels)\n    adv = _obj_rloo_advantages(rewards, group_size).unsqueeze(-1)\n    return _obj_clipped_pg_loss(logp, old_logp, adv, mask, clip_low, clip_high)\n\n\ndef reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):\n    logp = _obj_token_logprobs(logits, labels)\n    ref_logp = _obj_token_logprobs(ref_logits, labels)\n    denom = mask.sum()\n    if denom > 0:\n        pooled = (logp * mask).sum(0) / denom\n        pg = -(advantages * pooled).sum()\n        kl = logits.shape[0] * ((_obj_reverse_kl(logp, ref_logp) * mask).sum() / denom)\n    else:\n        pg = torch.zeros((), dtype=logp.dtype, device=logp.device)\n        kl = torch.zeros((), dtype=logp.dtype, device=logp.device)\n    return pg + beta * kl\n```"}

{"type": "text", "text": "agentId: a93426ba6b765ed85 (use SendMessage with to: 'a93426ba6b765ed85', summary: '<5-10 word recap>' to continue this agent)\n<usage>subagent_tokens: 63282\ntool_uses: 38\nduration_ms: 524362</usage>"}

from __future__ import annotations

import torch


# ---------------------------------------------------------------------------
# Masking / reduction primitives
# ---------------------------------------------------------------------------

def masked_mean(values, mask, axis=None):
    if axis is None:
        s = (values * mask).sum()
        d = mask.sum()
    else:
        s = (values * mask).sum(dim=axis)
        d = mask.sum(dim=axis)
    safe = d + (d == 0).to(d.dtype)
    return s / safe * (d != 0).to(s.dtype)


def masked_sum(values, mask, axis=None):
    if axis is None:
        return (values * mask).sum()
    return (values * mask).sum(dim=axis)


def masked_max(values, mask):
    neg_inf = torch.finfo(values.dtype).min
    masked_vals = torch.where(mask.bool(), values, torch.full_like(values, neg_inf))
    if mask.sum() == 0:
        return torch.zeros((), dtype=values.dtype, device=values.device)
    return masked_vals.max()


def _maskprim_safe_mean(values, mask):
    s = (values * mask).sum()
    d = mask.sum()
    safe = d + (d == 0).to(d.dtype)
    return s / safe * (d != 0).to(s.dtype)


def _maskprim_mvar(values, mask):
    mean = _maskprim_safe_mean(values, mask)
    centered = values - mean
    var = _maskprim_safe_mean(centered * centered, mask)
    return mean, var


def whiten(values, mask, shift_mean):
    mean, var = _maskprim_mvar(values, mask)
    whitened = (values - mean) / torch.sqrt(var + 1e-8)
    if not shift_mean:
        whitened = whitened + mean
    return whitened


def masked_whiten(values, mask, shift_mean):
    mean, var = _maskprim_mvar(values, mask)
    whitened = (values - mean) / torch.sqrt(var + 1e-8)
    whitened = whitened * mask
    if not shift_mean:
        whitened = whitened + mean
    return whitened


def advantage_mean_std(advantages, mask):
    mean, var = _maskprim_mvar(advantages, mask)
    std = torch.sqrt(var)
    return torch.stack([mean, std])


def normalize(x, eps):
    mean = x.mean()
    std = x.std(unbiased=True)
    return (x - mean) / (std + eps)


# ---------------------------------------------------------------------------
# Logprob / softmax primitives
# ---------------------------------------------------------------------------

def _lp_masked_mean(values, mask):
    denom = mask.sum()
    if denom == 0:
        return torch.zeros((), dtype=values.dtype, device=values.device)
    return (values * mask).sum() / denom


def logsumexp(x, axis):
    return torch.logsumexp(x, dim=axis)


def log_softmax(x, axis):
    return x - torch.logsumexp(x, dim=axis, keepdim=True)


def token_logprobs(logits, labels):
    lsm = torch.log_softmax(logits, dim=-1)
    return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)


def selective_logprobs(logits, labels, mask):
    lp = token_logprobs(logits, labels)
    return (lp * mask).sum(dim=-1)


def sequence_logprob(logits, labels, mask, length_normalize):
    lp = token_logprobs(logits, labels)
    total = (lp * mask).sum(dim=-1)
    if length_normalize:
        denom = mask.sum(dim=-1)
        safe = torch.where(denom > 0, denom, torch.ones_like(denom))
        out = total / safe
        return torch.where(denom > 0, out, torch.zeros_like(out))
    return total


def logprob_at_temperature(logits, labels, temperature):
    return token_logprobs(logits / temperature, labels)


def cross_entropy(logits, labels, ignore_index):
    labels = labels.long()
    flat_logits = logits.reshape(-1, logits.shape[-1])
    flat_labels = labels.reshape(-1)
    lsm = torch.log_softmax(flat_logits, dim=-1)
    valid = (flat_labels != ignore_index)
    safe_labels = torch.where(valid, flat_labels, torch.zeros_like(flat_labels))
    nll = -lsm.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)
    valid_f = valid.to(nll.dtype)
    denom = valid_f.sum()
    if denom == 0:
        return torch.zeros((), dtype=nll.dtype, device=nll.device)
    return (nll * valid_f).sum() / denom


def smoothed_nll(logits, labels, smoothing):
    lsm = torch.log_softmax(logits, dim=-1)
    nll = -lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)
    smooth = -lsm.mean(dim=-1)
    loss = (1.0 - smoothing) * nll + smoothing * smooth
    return loss.mean()


def entropy(logits, mask):
    lsm = torch.log_softmax(logits, dim=-1)
    p = torch.softmax(logits, dim=-1)
    ent = -(p * lsm).sum(dim=-1)
    return _lp_masked_mean(ent, mask)


def normalized_entropy(logits, mask):
    lsm = torch.log_softmax(logits, dim=-1)
    p = torch.softmax(logits, dim=-1)
    ent = -(p * lsm).sum(dim=-1)
    vocab = logits.shape[-1]
    norm = ent / torch.log(torch.tensor(float(vocab), dtype=ent.dtype, device=ent.device))
    return _lp_masked_mean(norm, mask)


# ---------------------------------------------------------------------------
# Tensor-op primitives
# ---------------------------------------------------------------------------

def top_p_mask(probs, p):
    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
    cumsum = torch.cumsum(sorted_probs, dim=-1)
    cumsum_before = cumsum - sorted_probs
    keep = cumsum_before < p
    keep[..., 0] = True
    kept = sorted_probs * keep.to(sorted_probs.dtype)
    kept = kept / kept.sum(dim=-1, keepdim=True)
    out = torch.zeros_like(kept)
    out.scatter_(-1, sorted_idx, kept)
    return out


def top_k_mask(logits, k):
    kth = torch.topk(logits, k, dim=-1).values[..., -1:]
    return logits >= kth


def argmax_tokens(logits):
    flipped = torch.flip(logits, dims=[-1])
    n = logits.shape[-1]
    idx = torch.argmax(flipped, dim=-1)
    return (n - 1) - idx


def mode_label(labels):
    def _mode_1d(v):
        uniq, counts = torch.unique(v, return_counts=True)
        max_count = counts.max()
        candidates = uniq[counts == max_count]
        return candidates.max()

    if labels.dim() == 1:
        return _mode_1d(labels)
    return torch.stack([_mode_1d(labels[i]) for i in range(labels.shape[0])])


def _top_positional_quantile(x, q, interpolation):
    xs, _ = torch.sort(x, dim=-1)
    n = xs.shape[0]
    pos = (n - 1) * q
    if interpolation == "lower":
        return xs[int(pos // 1)]
    lo = int(pos // 1)
    hi = min(lo + 1, n - 1)
    frac = pos - lo
    return xs[lo] * (1 - frac) + xs[hi] * frac


def median_reward(rewards):
    return _top_positional_quantile(rewards, 0.5, "linear")


def quantile_lower(x, q):
    return _top_positional_quantile(x, q, "lower")


def pad_mask_from_lengths(lengths, max_len):
    arange = torch.arange(max_len, device=lengths.device).unsqueeze(0)
    return (arange < lengths.unsqueeze(1)).to(torch.float32)


def first_nonzero_index(mask):
    nz = mask != 0
    any_nz = nz.any(dim=-1)
    idx = torch.argmax(nz.to(torch.long), dim=-1)
    return torch.where(any_nz, idx, torch.full_like(idx, -1))


def cumulative_max(x):
    out = torch.empty_like(x)
    out[0] = x[0]
    if x.shape[0] > 1:
        out[1:] = torch.cummax(x, dim=0)[0][:-1]
    return out


def bucketize_reward(rewards, boundaries):
    return torch.bucketize(rewards, boundaries, right=True)


# ---------------------------------------------------------------------------
# KL divergences / ratios
# ---------------------------------------------------------------------------

def kl_penalty(logp, ref_logp, estimator):
    logratio = logp - ref_logp
    if estimator == 'k1':
        return logratio
    elif estimator == 'k2':
        return 0.5 * logratio * logratio
    elif estimator == 'k3':
        return torch.exp(-logratio) + logratio - 1
    else:
        raise ValueError(estimator)


def reverse_kl(logp, ref_logp):
    logratio = logp - ref_logp
    return torch.exp(-logratio) + logratio - 1


def symmetric_kl(logp, ref_logp):
    logratio = logp - ref_logp
    reverse = torch.exp(-logratio) + logratio - 1
    forward = torch.exp(logratio) - logratio - 1
    return 0.5 * (reverse + forward)


def importance_ratio(logp, old_logp, clip):
    ratio = torch.exp(logp - old_logp)
    if clip is not None:
        ratio = torch.clamp(ratio, 1 - clip, 1 + clip)
    return ratio


def clip_fraction(logp, old_logp, clip):
    ratio = torch.exp(logp - old_logp)
    return ((ratio > 1 + clip) | (ratio < 1 - clip)).float().mean()


def bradley_terry_logit(chosen_reward, rejected_reward, beta):
    return beta * (chosen_reward - rejected_reward)


# ---------------------------------------------------------------------------
# RL advantages / returns / value losses
# ---------------------------------------------------------------------------

def _rl_masked_mean(x, mask):
    return (x * mask).sum() / mask.sum()


def grpo_advantages(rewards, group_size, scale_by_std):
    flat = rewards.reshape(-1)
    groups = flat.reshape(-1, group_size)
    mean = groups.mean(dim=1, keepdim=True)
    adv = groups - mean
    if scale_by_std:
        std = groups.std(dim=1, unbiased=False, keepdim=True)
        adv = adv / (std + 1e-6)
    return adv.reshape(-1)


def rloo_advantages(rewards, group_size):
    flat = rewards.reshape(-1)
    groups = flat.reshape(-1, group_size)
    group_sum = groups.sum(dim=1, keepdim=True)
    baseline = (group_sum - groups) / (group_size - 1)
    adv = groups - baseline
    return adv.reshape(-1)


def group_mean_baseline(rewards, group_size):
    flat = rewards.reshape(-1)
    groups = flat.reshape(-1, group_size)
    mean = groups.mean(dim=1, keepdim=True)
    adv = groups - mean
    return adv.reshape(-1)


def gae(rewards, values, next_value, gamma, lam):
    T = rewards.shape[0]
    if not torch.is_tensor(next_value):
        next_value = torch.tensor(next_value, dtype=values.dtype)
    adv = torch.zeros_like(rewards)
    gae_acc = torch.zeros_like(rewards[0])
    next_v = next_value
    for t in range(T - 1, -1, -1):
        delta = rewards[t] + gamma * next_v - values[t]
        gae_acc = delta + gamma * lam * gae_acc
        adv[t] = gae_acc
        next_v = values[t]
    return adv


def discounted_returns(rewards, gamma):
    T = rewards.shape[0]
    out = torch.zeros_like(rewards)
    acc = torch.zeros_like(rewards[0])
    for t in range(T - 1, -1, -1):
        acc = rewards[t] + gamma * acc
        out[t] = acc
    return out


def lambda_returns(rewards, values, next_value, gamma, lam):
    T = rewards.shape[0]
    if not torch.is_tensor(next_value):
        next_value = torch.tensor(next_value, dtype=values.dtype)
    out = torch.zeros_like(rewards)
    next_ret = next_value
    for t in range(T - 1, -1, -1):
        next_v = next_value if t == T - 1 else values[t + 1]
        next_ret = rewards[t] + gamma * ((1 - lam) * next_v + lam * next_ret)
        out[t] = next_ret
    return out


def clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):
    ratio = torch.exp(logp - old_logp)
    unclipped = ratio * advantages
    clipped = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) * advantages
    per = -torch.min(unclipped, clipped)
    return _rl_masked_mean(per, mask)


def value_loss(values, old_values, returns, clip):
    unclipped = (values - returns) ** 2
    vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)
    clipped = (vpred_clipped - returns) ** 2
    return 0.5 * torch.max(unclipped, clipped).mean()


def huber_value_loss(values, returns, delta):
    x = values - returns
    absx = x.abs()
    huber = torch.where(absx <= delta, 0.5 * x ** 2, delta * (absx - 0.5 * delta))
    return huber.mean()


# ---------------------------------------------------------------------------
# Preference losses
# ---------------------------------------------------------------------------

def dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
    logits = beta * ((pc - rc) - (pr - rr))
    loss = (
        -(1 - label_smoothing) * torch.nn.functional.logsigmoid(logits)
        - label_smoothing * torch.nn.functional.logsigmoid(-logits)
    )
    return loss.mean()


def ipo_loss(pc, pr, rc, rr, beta):
    h = (pc - rc) - (pr - rr)
    loss = (h - 1 / (2 * beta)) ** 2
    return loss.mean()


# ---------------------------------------------------------------------------
# Composite objectives
# ---------------------------------------------------------------------------

def _obj_token_logprobs(logits, labels):
    lsm = torch.log_softmax(logits, dim=-1)
    return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)


def _obj_sequence_logprob(logits, labels, mask, length_normalize=False):
    lp = _obj_token_logprobs(logits, labels)
    s = (lp * mask).sum(-1)
    if length_normalize:
        s = s / mask.sum(-1).clamp(min=1.0)
    return s


def _obj_masked_mean(values, mask):
    s = mask.sum()
    if s > 0:
        return (values * mask).sum() / s
    return torch.zeros((), dtype=values.dtype, device=values.device)


def _obj_clipped_pg_loss(logp, old_logp, adv, mask, clip_low, clip_high):
    ratio = torch.exp(logp - old_logp)
    unclipped = ratio * adv
    clipped = torch.clamp(ratio, 1.0 - clip_low, 1.0 + clip_high) * adv
    per = -torch.min(unclipped, clipped)
    return _obj_masked_mean(per, mask)


def _obj_grpo_advantages(rewards, group_size, scale_by_std):
    flat = rewards.reshape(-1)
    groups = flat.reshape(-1, group_size)
    adv = groups - groups.mean(1, keepdim=True)
    if scale_by_std:
        adv = adv / (groups.std(1, unbiased=False, keepdim=True) + 1e-6)
    return adv.reshape(-1)


def _obj_rloo_advantages(rewards, group_size):
    flat = rewards.reshape(-1)
    groups = flat.reshape(-1, group_size)
    gsum = groups.sum(1, keepdim=True)
    baseline = (gsum - groups) / (group_size - 1)
    return (groups - baseline).reshape(-1)


def _obj_kl_penalty(logp, ref_logp, estimator):
    logratio = logp - ref_logp
    if estimator == "k1":
        return logratio
    if estimator == "k2":
        return 0.5 * logratio ** 2
    return torch.exp(-logratio) + logratio - 1.0


def _obj_reverse_kl(logp, ref_logp):
    lr = logp - ref_logp
    return torch.exp(-lr) + lr - 1.0


def _obj_dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
    logits = beta * ((pc - rc) - (pr - rr))
    loss = (-(1.0 - label_smoothing) * torch.nn.functional.logsigmoid(logits)
            - label_smoothing * torch.nn.functional.logsigmoid(-logits))
    return loss.mean()


def _obj_gae(rewards, values, next_value, gamma, lam):
    T = rewards.shape[0]
    if not torch.is_tensor(next_value):
        next_value = torch.tensor(next_value, dtype=values.dtype)
    adv = torch.zeros_like(rewards)
    last = torch.zeros_like(rewards[0])
    next_v = next_value
    for t in range(T - 1, -1, -1):
        delta = rewards[t] + gamma * next_v - values[t]
        last = delta + gamma * lam * last
        adv[t] = last
        next_v = values[t]
    return adv


def _obj_whiten(x, eps=1e-8):
    return (x - x.mean()) / (x.std(unbiased=False) + eps)


def _obj_value_loss(values, old_values, returns, clip):
    unclipped = (values - returns) ** 2
    vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)
    clipped = (vpred_clipped - returns) ** 2
    return 0.5 * torch.max(unclipped, clipped).mean()


def dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,
                      chosen_labels, rejected_labels, chosen_mask, rejected_mask,
                      beta, label_smoothing):
    pc = _obj_sequence_logprob(pc_logits, chosen_labels, chosen_mask, length_normalize=False)
    pr = _obj_sequence_logprob(pr_logits, rejected_labels, rejected_mask, length_normalize=False)
    rc = _obj_sequence_logprob(rc_logits, chosen_labels, chosen_mask, length_normalize=False)
    rr = _obj_sequence_logprob(rr_logits, rejected_labels, rejected_mask, length_normalize=False)
    return _obj_dpo_loss(pc, pr, rc, rr, beta, label_smoothing)


def grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,
                   rewards, group_size, beta, clip_low, clip_high, scale_by_std,
                   kl_estimator):
    logp = _obj_token_logprobs(logits, labels)
    old_logp = _obj_token_logprobs(old_logits, labels)
    ref_logp = _obj_token_logprobs(ref_logits, labels)
    adv = _obj_grpo_advantages(rewards, group_size, scale_by_std).unsqueeze(-1)
    pg_loss = _obj_clipped_pg_loss(logp, old_logp, adv, completion_mask, clip_low, clip_high)
    kl = _obj_kl_penalty(logp, ref_logp, kl_estimator)
    return pg_loss + beta * _obj_masked_mean(kl, completion_mask)


def ppo_objective(rewards, values, old_values, logp, old_logp, next_value,
                  gamma, lam, clip_low, clip_high, vf_clip, vf_coef):
    advantages = _obj_gae(rewards, values, next_value, gamma, lam)
    returns = advantages + values
    advantages = _obj_whiten(advantages)
    mask = torch.ones_like(logp)
    pg_loss = _obj_clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high)
    vf = _obj_value_loss(values, old_values, returns, vf_clip)
    return pg_loss + vf_coef * vf


def rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):
    logp = _obj_token_logprobs(logits, labels)
    old_logp = _obj_token_logprobs(old_logits, labels)
    adv = _obj_rloo_advantages(rewards, group_size).unsqueeze(-1)
    return _obj_clipped_pg_loss(logp, old_logp, adv, mask, clip_low, clip_high)


def reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):
    logp = _obj_token_logprobs(logits, labels)
    ref_logp = _obj_token_logprobs(ref_logits, labels)
    denom = mask.sum()
    if denom > 0:
        pooled = (logp * mask).sum(0) / denom
        pg = -(advantages * pooled).sum()
        kl = logits.shape[0] * ((_obj_reverse_kl(logp, ref_logp) * mask).sum() / denom)
    else:
        pg = torch.zeros((), dtype=logp.dtype, device=logp.device)
        kl = torch.zeros((), dtype=logp.dtype, device=logp.device)
    return pg + beta * kl

from __future__ import annotations

import torch


# ---------------------------------------------------------------------------
# Masking / reduction primitives
# ---------------------------------------------------------------------------

def masked_mean(values, mask, axis=None):
    if axis is None:
        s = (values * mask).sum()
        d = mask.sum()
    else:
        s = (values * mask).sum(dim=axis)
        d = mask.sum(dim=axis)
    safe = d + (d == 0).to(d.dtype)
    return s / safe * (d != 0).to(s.dtype)


def masked_sum(values, mask, axis=None):
    if axis is None:
        return (values * mask).sum()
    return (values * mask).sum(dim=axis)


def masked_max(values, mask):
    neg_inf = torch.finfo(values.dtype).min
    masked_vals = torch.where(mask.bool(), values, torch.full_like(values, neg_inf))
    if mask.sum() == 0:
        return torch.zeros((), dtype=values.dtype, device=values.device)
    return masked_vals.max()


def _maskprim_safe_mean(values, mask):
    s = (values * mask).sum()
    d = mask.sum()
    safe = d + (d == 0).to(d.dtype)
    return s / safe * (d != 0).to(s.dtype)


def _maskprim_mvar(values, mask):
    mean = _maskprim_safe_mean(values, mask)
    centered = values - mean
    var = _maskprim_safe_mean(centered * centered, mask)
    return mean, var


def whiten(values, mask, shift_mean):
    mean, var = _maskprim_mvar(values, mask)
    whitened = (values - mean) / torch.sqrt(var + 1e-8)
    if not shift_mean:
        whitened = whitened + mean
    return whitened


def masked_whiten(values, mask, shift_mean):
    mean, var = _maskprim_mvar(values, mask)
    whitened = (values - mean) / torch.sqrt(var + 1e-8)
    whitened = whitened * mask
    if not shift_mean:
        whitened = whitened + mean
    return whitened


def advantage_mean_std(advantages, mask):
    mean, var = _maskprim_mvar(advantages, mask)
    std = torch.sqrt(var)
    return torch.stack([mean, std])


def normalize(x, eps):
    mean = x.mean()
    std = x.std(unbiased=True)
    return (x - mean) / (std + eps)


# ---------------------------------------------------------------------------
# Logprob / softmax primitives
# ---------------------------------------------------------------------------

def _lp_masked_mean(values, mask):
    denom = mask.sum()
    if denom == 0:
        return torch.zeros((), dtype=values.dtype, device=values.device)
    return (values * mask).sum() / denom


def logsumexp(x, axis):
    return torch.logsumexp(x, dim=axis)


def log_softmax(x, axis):
    return x - torch.logsumexp(x, dim=axis, keepdim=True)


def token_logprobs(logits, labels):
    lsm = torch.log_softmax(logits, dim=-1)
    return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)


def selective_logprobs(logits, labels, mask):
    lp = token_logprobs(logits, labels)
    return (lp * mask).sum(dim=-1)


def sequence_logprob(logits, labels, mask, length_normalize):
    lp = token_logprobs(logits, labels)
    total = (lp * mask).sum(dim=-1)
    if length_normalize:
        denom = mask.sum(dim=-1)
        safe = torch.where(denom > 0, denom, torch.ones_like(denom))
        out = total / safe
        return torch.where(denom > 0, out, torch.zeros_like(out))
    return total


def logprob_at_temperature(logits, labels, temperature):
    return token_logprobs(logits / temperature, labels)


def cross_entropy(logits, labels, ignore_index):
    labels = labels.long()
    flat_logits = logits.reshape(-1, logits.shape[-1])
    flat_labels = labels.reshape(-1)
    lsm = torch.log_softmax(flat_logits, dim=-1)
    valid = (flat_labels != ignore_index)
    safe_labels = torch.where(valid, flat_labels, torch.zeros_like(flat_labels))
    nll = -lsm.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)
    valid_f = valid.to(nll.dtype)
    denom = valid_f.sum()
    if denom == 0:
        return torch.zeros((), dtype=nll.dtype, device=nll.device)
    return (nll * valid_f).sum() / denom


def smoothed_nll(logits, labels, smoothing):
    lsm = torch.log_softmax(logits, dim=-1)
    nll = -lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)
    smooth = -lsm.mean(dim=-1)
    loss = (1.0 - smoothing) * nll + smoothing * smooth
    return loss.mean()


def entropy(logits, mask):
    lsm = torch.log_softmax(logits, dim=-1)
    p = torch.softmax(logits, dim=-1)
    ent = -(p * lsm).sum(dim=-1)
    return _lp_masked_mean(ent, mask)


def normalized_entropy(logits, mask):
    lsm = torch.log_softmax(logits, dim=-1)
    p = torch.softmax(logits, dim=-1)
    ent = -(p * lsm).sum(dim=-1)
    vocab = logits.shape[-1]
    norm = ent / torch.log(torch.tensor(float(vocab), dtype=ent.dtype, device=ent.device))
    return _lp_masked_mean(norm, mask)


# ---------------------------------------------------------------------------
# Tensor-op primitives
# ---------------------------------------------------------------------------

def top_p_mask(probs, p):
    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
    cumsum = torch.cumsum(sorted_probs, dim=-1)
    cumsum_before = cumsum - sorted_probs
    keep = cumsum_before < p
    keep[..., 0] = True
    kept = sorted_probs * keep.to(sorted_probs.dtype)
    kept = kept / kept.sum(dim=-1, keepdim=True)
    out = torch.zeros_like(kept)
    out.scatter_(-1, sorted_idx, kept)
    return out


def top_k_mask(logits, k):
    kth = torch.topk(logits, k, dim=-1).values[..., -1:]
    return logits >= kth


def argmax_tokens(logits):
    flipped = torch.flip(logits, dims=[-1])
    n = logits.shape[-1]
    idx = torch.argmax(flipped, dim=-1)
    return (n - 1) - idx


def mode_label(labels):
    def _mode_1d(v):
        uniq, counts = torch.unique(v, return_counts=True)
        max_count = counts.max()
        candidates = uniq[counts == max_count]
        return candidates.max()

    if labels.dim() == 1:
        return _mode_1d(labels)
    return torch.stack([_mode_1d(labels[i]) for i in range(labels.shape[0])])


def _top_positional_quantile(x, q, interpolation):
    xs, _ = torch.sort(x, dim=-1)
    n = xs.shape[0]
    pos = (n - 1) * q
    if interpolation == "lower":
        return xs[int(pos // 1)]
    lo = int(pos // 1)
    hi = min(lo + 1, n - 1)
    frac = pos - lo
    return xs[lo] * (1 - frac) + xs[hi] * frac


def median_reward(rewards):
    return _top_positional_quantile(rewards, 0.5, "linear")


def quantile_lower(x, q):
    return _top_positional_quantile(x, q, "lower")


def pad_mask_from_lengths(lengths, max_len):
    arange = torch.arange(max_len, device=lengths.device).unsqueeze(0)
    return (arange < lengths.unsqueeze(1)).to(torch.float32)


def first_nonzero_index(mask):
    nz = mask != 0
    any_nz = nz.any(dim=-1)
    idx = torch.argmax(nz.to(torch.long), dim=-1)
    return torch.where(any_nz, idx, torch.full_like(idx, -1))


def cumulative_max(x):
    out = torch.empty_like(x)
    out[0] = x[0]
    if x.shape[0] > 1:
        out[1:] = torch.cummax(x, dim=0)[0][:-1]
    return out


def bucketize_reward(rewards, boundaries):
    return torch.bucketize(rewards, boundaries, right=True)


# ---------------------------------------------------------------------------
# KL divergences / ratios
# ---------------------------------------------------------------------------

def kl_penalty(logp, ref_logp, estimator):
    logratio = logp - ref_logp
    if estimator == 'k1':
        return logratio
    elif estimator == 'k2':
        return 0.5 * logratio * logratio
    elif estimator == 'k3':
        return torch.exp(-logratio) + logratio - 1
    else:
        raise ValueError(estimator)


def reverse_kl(logp, ref_logp):
    logratio = logp - ref_logp
    return torch.exp(-logratio) + logratio - 1


def symmetric_kl(logp, ref_logp):
    logratio = logp - ref_logp
    reverse = torch.exp(-logratio) + logratio - 1
    forward = torch.exp(logratio) - logratio - 1
    return 0.5 * (reverse + forward)


def importance_ratio(logp, old_logp, clip):
    ratio = torch.exp(logp - old_logp)
    if clip is not None:
        ratio = torch.clamp(ratio, 1 - clip, 1 + clip)
    return ratio


def clip_fraction(logp, old_logp, clip):
    ratio = torch.exp(logp - old_logp)
    return ((ratio > 1 + clip) | (ratio < 1 - clip)).float().mean()


def bradley_terry_logit(chosen_reward, rejected_reward, beta):
    return beta * (chosen_reward - rejected_reward)


# ---------------------------------------------------------------------------
# RL advantages / returns / value losses
# ---------------------------------------------------------------------------

def _rl_masked_mean(x, mask):
    return (x * mask).sum() / mask.sum()


def grpo_advantages(rewards, group_size, scale_by_std):
    flat = rewards.reshape(-1)
    groups = flat.reshape(-1, group_size)
    mean = groups.mean(dim=1, keepdim=True)
    adv = groups - mean
    if scale_by_std:
        std = groups.std(dim=1, unbiased=False, keepdim=True)
        adv = adv / (std + 1e-6)
    return adv.reshape(-1)


def rloo_advantages(rewards, group_size):
    flat = rewards.reshape(-1)
    groups = flat.reshape(-1, group_size)
    group_sum = groups.sum(dim=1, keepdim=True)
    baseline = (group_sum - groups) / (group_size - 1)
    adv = groups - baseline
    return adv.reshape(-1)


def group_mean_baseline(rewards, group_size):
    flat = rewards.reshape(-1)
    groups = flat.reshape(-1, group_size)
    mean = groups.mean(dim=1, keepdim=True)
    adv = groups - mean
    return adv.reshape(-1)


def gae(rewards, values, next_value, gamma, lam):
    T = rewards.shape[0]
    if not torch.is_tensor(next_value):
        next_value = torch.tensor(next_value, dtype=values.dtype)
    adv = torch.zeros_like(rewards)
    gae_acc = torch.zeros_like(rewards[0])
    next_v = next_value
    for t in range(T - 1, -1, -1):
        delta = rewards[t] + gamma * next_v - values[t]
        gae_acc = delta + gamma * lam * gae_acc
        adv[t] = gae_acc
        next_v = values[t]
    return adv


def discounted_returns(rewards, gamma):
    T = rewards.shape[0]
    out = torch.zeros_like(rewards)
    acc = torch.zeros_like(rewards[0])
    for t in range(T - 1, -1, -1):
        acc = rewards[t] + gamma * acc
        out[t] = acc
    return out


def lambda_returns(rewards, values, next_value, gamma, lam):
    T = rewards.shape[0]
    if not torch.is_tensor(next_value):
        next_value = torch.tensor(next_value, dtype=values.dtype)
    out = torch.zeros_like(rewards)
    next_ret = next_value
    for t in range(T - 1, -1, -1):
        next_v = next_value if t == T - 1 else values[t + 1]
        next_ret = rewards[t] + gamma * ((1 - lam) * next_v + lam * next_ret)
        out[t] = next_ret
    return out


def clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):
    ratio = torch.exp(logp - old_logp)
    unclipped = ratio * advantages
    clipped = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) * advantages
    per = -torch.min(unclipped, clipped)
    return _rl_masked_mean(per, mask)


def value_loss(values, old_values, returns, clip):
    unclipped = (values - returns) ** 2
    vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)
    clipped = (vpred_clipped - returns) ** 2
    return 0.5 * torch.max(unclipped, clipped).mean()


def huber_value_loss(values, returns, delta):
    x = values - returns
    absx = x.abs()
    huber = torch.where(absx <= delta, 0.5 * x ** 2, delta * (absx - 0.5 * delta))
    return huber.mean()


# ---------------------------------------------------------------------------
# Preference losses
# ---------------------------------------------------------------------------

def dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
    logits = beta * ((pc - rc) - (pr - rr))
    loss = (
        -(1 - label_smoothing) * torch.nn.functional.logsigmoid(logits)
        - label_smoothing * torch.nn.functional.logsigmoid(-logits)
    )
    return loss.mean()


def ipo_loss(pc, pr, rc, rr, beta):
    h = (pc - rc) - (pr - rr)
    loss = (h - 1 / (2 * beta)) ** 2
    return loss.mean()


# ---------------------------------------------------------------------------
# Composite objectives
# ---------------------------------------------------------------------------

def _obj_token_logprobs(logits, labels):
    lsm = torch.log_softmax(logits, dim=-1)
    return lsm.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)


def _obj_sequence_logprob(logits, labels, mask, length_normalize=False):
    lp = _obj_token_logprobs(logits, labels)
    s = (lp * mask).sum(-1)
    if length_normalize:
        s = s / mask.sum(-1).clamp(min=1.0)
    return s


def _obj_masked_mean(values, mask):
    s = mask.sum()
    if s > 0:
        return (values * mask).sum() / s
    return torch.zeros((), dtype=values.dtype, device=values.device)


def _obj_clipped_pg_loss(logp, old_logp, adv, mask, clip_low, clip_high):
    ratio = torch.exp(logp - old_logp)
    unclipped = ratio * adv
    clipped = torch.clamp(ratio, 1.0 - clip_low, 1.0 + clip_high) * adv
    per = -torch.min(unclipped, clipped)
    return _obj_masked_mean(per, mask)


def _obj_grpo_advantages(rewards, group_size, scale_by_std):
    flat = rewards.reshape(-1)
    groups = flat.reshape(-1, group_size)
    adv = groups - groups.mean(1, keepdim=True)
    if scale_by_std:
        adv = adv / (groups.std(1, unbiased=False, keepdim=True) + 1e-6)
    return adv.reshape(-1)


def _obj_rloo_advantages(rewards, group_size):
    flat = rewards.reshape(-1)
    groups = flat.reshape(-1, group_size)
    gsum = groups.sum(1, keepdim=True)
    baseline = (gsum - groups) / (group_size - 1)
    return (groups - baseline).reshape(-1)


def _obj_kl_penalty(logp, ref_logp, estimator):
    logratio = logp - ref_logp
    if estimator == "k1":
        return logratio
    if estimator == "k2":
        return 0.5 * logratio ** 2
    return torch.exp(-logratio) + logratio - 1.0


def _obj_reverse_kl(logp, ref_logp):
    lr = logp - ref_logp
    return torch.exp(-lr) + lr - 1.0


def _obj_dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
    logits = beta * ((pc - rc) - (pr - rr))
    loss = (-(1.0 - label_smoothing) * torch.nn.functional.logsigmoid(logits)
            - label_smoothing * torch.nn.functional.logsigmoid(-logits))
    return loss.mean()


def _obj_gae(rewards, values, next_value, gamma, lam):
    T = rewards.shape[0]
    if not torch.is_tensor(next_value):
        next_value = torch.tensor(next_value, dtype=values.dtype)
    adv = torch.zeros_like(rewards)
    last = torch.zeros_like(rewards[0])
    next_v = next_value
    for t in range(T - 1, -1, -1):
        delta = rewards[t] + gamma * next_v - values[t]
        last = delta + gamma * lam * last
        adv[t] = last
        next_v = values[t]
    return adv


def _obj_whiten(x, eps=1e-8):
    return (x - x.mean()) / (x.std(unbiased=False) + eps)


def _obj_value_loss(values, old_values, returns, clip):
    unclipped = (values - returns) ** 2
    vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)
    clipped = (vpred_clipped - returns) ** 2
    return 0.5 * torch.max(unclipped, clipped).mean()


def dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,
                      chosen_labels, rejected_labels, chosen_mask, rejected_mask,
                      beta, label_smoothing):
    pc = _obj_sequence_logprob(pc_logits, chosen_labels, chosen_mask, length_normalize=False)
    pr = _obj_sequence_logprob(pr_logits, rejected_labels, rejected_mask, length_normalize=False)
    rc = _obj_sequence_logprob(rc_logits, chosen_labels, chosen_mask, length_normalize=False)
    rr = _obj_sequence_logprob(rr_logits, rejected_labels, rejected_mask, length_normalize=False)
    return _obj_dpo_loss(pc, pr, rc, rr, beta, label_smoothing)


def grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,
                   rewards, group_size, beta, clip_low, clip_high, scale_by_std,
                   kl_estimator):
    logp = _obj_token_logprobs(logits, labels)
    old_logp = _obj_token_logprobs(old_logits, labels)
    ref_logp = _obj_token_logprobs(ref_logits, labels)
    adv = _obj_grpo_advantages(rewards, group_size, scale_by_std).unsqueeze(-1)
    pg_loss = _obj_clipped_pg_loss(logp, old_logp, adv, completion_mask, clip_low, clip_high)
    kl = _obj_kl_penalty(logp, ref_logp, kl_estimator)
    return pg_loss + beta * _obj_masked_mean(kl, completion_mask)


def ppo_objective(rewards, values, old_values, logp, old_logp, next_value,
                  gamma, lam, clip_low, clip_high, vf_clip, vf_coef):
    advantages = _obj_gae(rewards, values, next_value, gamma, lam)
    returns = advantages + values
    advantages = _obj_whiten(advantages)
    mask = torch.ones_like(logp)
    pg_loss = _obj_clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high)
    vf = _obj_value_loss(values, old_values, returns, vf_clip)
    return pg_loss + vf_coef * vf


def rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):
    logp = _obj_token_logprobs(logits, labels)
    old_logp = _obj_token_logprobs(old_logits, labels)
    adv = _obj_rloo_advantages(rewards, group_size).unsqueeze(-1)
    return _obj_clipped_pg_loss(logp, old_logp, adv, mask, clip_low, clip_high)


def reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):
    logp = _obj_token_logprobs(logits, labels)
    ref_logp = _obj_token_logprobs(ref_logits, labels)
    denom = mask.sum()
    if denom > 0:
        pooled = (logp * mask).sum(0) / denom
        pg = -(advantages * pooled).sum()
        kl = logits.shape[0] * ((_obj_reverse_kl(logp, ref_logp) * mask).sum() / denom)
    else:
        pg = torch.zeros((), dtype=logp.dtype, device=logp.device)
        kl = torch.zeros((), dtype=logp.dtype, device=logp.device)
    return pg + beta * kl

import torch, math, traceback
import reforacle as R
import objectives as O

g = torch.Generator().manual_seed(12345)
def rn(*s): return torch.randn(*s, generator=g)
def ru(*s): return torch.rand(*s, generator=g)
def mask(*s): return (torch.rand(*s, generator=g) > 0.4).float()

ATOL_P, RTOL_P = 1e-6, 1e-5
ATOL_C, RTOL_C = 1e-5, 1e-4

fails = {}
def chk(name, a, b, atol, rtol):
    try:
        if isinstance(a, tuple): a = a[0]
        if isinstance(b, tuple): b = b[0]
        a = a if torch.is_tensor(a) else torch.tensor(a)
        b = b if torch.is_tensor(b) else torch.tensor(b)
        if a.shape != b.shape:
            fails.setdefault(name, []).append(f"shape {a.shape} vs {b.shape}")
            return
        if a.dtype in (torch.long, torch.bool) or b.dtype in (torch.long, torch.bool):
            if not torch.equal(a, b):
                fails.setdefault(name, []).append(f"int/bool mismatch\n{a}\n{b}")
            return
        if not torch.allclose(a.float(), b.float(), atol=atol, rtol=rtol, equal_nan=True):
            d = (a.float()-b.float()).abs().max().item()
            fails.setdefault(name, []).append(f"maxdiff {d}")
    except Exception as e:
        fails.setdefault(name, []).append(f"EXC {e}")

def run(name, fn_args, atol=ATOL_P, rtol=RTOL_P):
    try:
        r = getattr(R, name)(*[a.clone() if torch.is_tensor(a) else a for a in fn_args])
    except Exception as e:
        # oracle itself errors -> skip this case
        return
    try:
        o = getattr(O, name)(*[a.clone() if torch.is_tensor(a) else a for a in fn_args])
    except Exception as e:
        fails.setdefault(name, []).append(f"OURS EXC {e}\n{traceback.format_exc()}")
        return
    chk(name, o, r, atol, rtol)

for trial in range(40):
    N, T, V = 6, 5, 7
    logits = rn(N, T, V); labels = torch.randint(0, V, (N, T), generator=g)
    m = mask(N, T)
    # masking primitives
    v1 = rn(8); mk1 = mask(8)
    run("masked_mean", [v1, mk1]); run("masked_mean", [rn(4,5), mask(4,5), 0]); run("masked_mean", [rn(4,5), mask(4,5), 1])
    run("masked_sum", [v1, mk1]); run("masked_sum", [rn(4,5), mask(4,5), 1])
    run("masked_max", [v1, mk1])
    run("whiten", [v1, mk1, True]); run("whiten", [v1, mk1, False])
    run("masked_whiten", [v1, mk1, True]); run("masked_whiten", [v1, mk1, False])
    run("advantage_mean_std", [v1, mk1])
    run("normalize", [v1, 1e-8]); run("normalize", [rn(10), 1e-5])
    # logprob
    run("logsumexp", [logits, -1]); run("logsumexp", [logits, 1])
    run("log_softmax", [logits, -1]); run("log_softmax", [logits, 0])
    run("token_logprobs", [logits, labels])
    run("selective_logprobs", [logits, labels, m])
    run("sequence_logprob", [logits, labels, m, True]); run("sequence_logprob", [logits, labels, m, False])
    run("logprob_at_temperature", [logits, labels, 0.7]); run("logprob_at_temperature", [logits, labels, 2.0])
    run("cross_entropy", [rn(N,V), torch.randint(0,V,(N,),generator=g), -100])
    lab_ig = labels.clone(); lab_ig[0,0]=-100
    run("cross_entropy", [logits, lab_ig, -100])
    run("smoothed_nll", [logits, labels, 0.1]); run("smoothed_nll", [logits, labels, 0.0])
    run("entropy", [logits, m]); run("normalized_entropy", [logits, m])
    # tensor ops
    probs = torch.softmax(rn(4,9), dim=-1)
    run("top_p_mask", [probs, 0.8]); run("top_p_mask", [probs, 0.3]); run("top_p_mask", [probs, 0.95])
    run("top_k_mask", [rn(4,9), 3]); run("top_k_mask", [rn(4,9), 1])
    run("argmax_tokens", [rn(3,4,6)])
    run("mode_label", [torch.randint(0,4,(10,),generator=g)]); run("mode_label", [torch.randint(0,4,(5,8),generator=g)])
    run("median_reward", [rn(9)]); run("median_reward", [rn(10)])
    run("quantile_lower", [rn(11), 0.3]); run("quantile_lower", [rn(11), 0.75])
    run("pad_mask_from_lengths", [torch.randint(0,6,(7,),generator=g), 6])
    run("first_nonzero_index", [mask(5,8)]); run("first_nonzero_index", [torch.zeros(3,4)])
    run("cumulative_max", [rn(10)]); run("cumulative_max", [rn(6,3)])
    run("bucketize_reward", [rn(10), torch.tensor([-1.0,0.0,1.0])])
    # KL
    lp = rn(6,4); rlp = rn(6,4); olp = lp + 0.05*rn(6,4)
    for est in ["k1","k2","k3"]:
        run("kl_penalty", [lp, rlp, est])
    run("reverse_kl", [lp, rlp]); run("symmetric_kl", [lp, rlp])
    run("importance_ratio", [lp, olp, None]); run("importance_ratio", [lp, olp, 0.2])
    run("clip_fraction", [lp, olp, 0.1])
    run("bradley_terry_logit", [rn(5), rn(5), 0.5])
    # RL
    rew = rn(12)
    run("grpo_advantages", [rew, 4, True]); run("grpo_advantages", [rew, 4, False]); run("grpo_advantages", [rew, 3, True])
    run("rloo_advantages", [rew, 4]); run("group_mean_baseline", [rew, 4])
    rw=rn(5); vl=rn(5); nv=rn(())
    run("gae", [rw, vl, nv, 0.99, 0.95]); run("gae", [rn(5,3), rn(5,3), rn(3), 0.9, 0.8])
    run("discounted_returns", [rw, 0.97])
    run("lambda_returns", [rw, vl, nv, 0.99, 0.95])
    lg=rn(6,5); olg=lg+0.05*rn(6,5); adv=rn(6,5); mm=mask(6,5)
    run("clipped_pg_loss", [lg, olg, adv, mm, 0.2, 0.2]); run("clipped_pg_loss", [lg, olg, adv, mm, 0.1, 0.3])
    run("value_loss", [rn(8), rn(8), rn(8), 0.2])
    run("huber_value_loss", [rn(8), rn(8), 1.0]); run("huber_value_loss", [rn(8), rn(8), 0.5])
    # preference
    pc,pr2,rc,rr=rn(5),rn(5),rn(5),rn(5)
    run("dpo_loss", [pc,pr2,rc,rr,0.1,0.0]); run("dpo_loss", [pc,pr2,rc,rr,0.1,0.1])
    run("ipo_loss", [pc,pr2,rc,rr,0.5])
    # composite
    lg2=rn(N,T,V); olg2=lg2+0.05*rn(N,T,V); rlg2=lg2+0.1*rn(N,T,V)
    run("dpo_sequence_loss", [rn(N,T,V),rn(N,T,V),rn(N,T,V),rn(N,T,V),
                              torch.randint(0,V,(N,T),generator=g),torch.randint(0,V,(N,T),generator=g),
                              mask(N,T),mask(N,T),0.1,0.0], ATOL_C, RTOL_C)
    for est in ["k1","k2","k3"]:
        run("grpo_objective", [lg2,olg2,rlg2,labels,m,rn(N),3,0.05,0.2,0.2,True,est], ATOL_C, RTOL_C)
    run("grpo_objective", [lg2,olg2,rlg2,labels,m,rn(N),2,0.0,0.2,0.2,False,"k1"], ATOL_C, RTOL_C)
    run("ppo_objective", [rn(7),rn(7),rn(7),rn(7),rn(7)+0.05*rn(7),rn(()),0.99,0.95,0.2,0.2,0.2,0.5], ATOL_C, RTOL_C)
    run("rloo_objective", [lg2,olg2,labels,m,rn(N),3,0.2,0.2], ATOL_C, RTOL_C)
    run("reverse_kl_objective", [lg2,rlg2,labels,m,rn(N,T),0.1], ATOL_C, RTOL_C)

allfns = [n for n in dir(O) if not n.startswith("_") and callable(getattr(O,n)) and n not in ("torch","annotations")]
tested = set()
import inspect
# report
print("=== FAILURES ===")
if not fails:
    print("NONE")
for k,v in fails.items():
    print(f"\n{k}: {len(v)} fails")
    print("  ", v[0][:300])
print("\n=== functions with zero failures across trials ===")
target = [f for f in dir(R) if not f.startswith("_")]
print("total oracle fns probed OK:", len([f for f in target if f in [n for n in dir(O)] and f not in fails]))
print("FAILING FUNCS:", sorted(fails.keys()))

import torch, traceback
import reforacle as R
import objectives as O

g = torch.Generator().manual_seed(999)
fails = []
def run(name, args, atol=1e-5, rtol=1e-4):
    try:
        r = getattr(R, name)(*[a.clone() if torch.is_tensor(a) else a for a in args])
    except Exception as e:
        return  # oracle errors -> not a graded case
    try:
        o = getattr(O, name)(*[a.clone() if torch.is_tensor(a) else a for a in args])
    except Exception as e:
        fails.append(f"{name}: OURS EXC {e}"); return
    a = r if torch.is_tensor(r) else torch.tensor(r)
    b = o if torch.is_tensor(o) else torch.tensor(o)
    if a.shape != b.shape:
        fails.append(f"{name}: shape {b.shape} vs oracle {a.shape}"); return
    if a.dtype != b.dtype:
        fails.append(f"{name}: DTYPE {b.dtype} vs oracle {a.dtype}")
    if a.dtype in (torch.long, torch.bool):
        if not torch.equal(a,b): fails.append(f"{name}: int mismatch {a} vs {b}")
        return
    if not torch.allclose(a.float(), b.float(), atol=atol, rtol=rtol, equal_nan=True):
        fails.append(f"{name}: maxdiff {(a.float()-b.float()).abs().max().item()}")

# float64 dtype preservation
d = torch.float64
run("masked_mean", [torch.randn(6,dtype=d), (torch.rand(6)>0.3).double()])
run("normalize", [torch.randn(6,dtype=d), 1e-8])
run("whiten", [torch.randn(6,dtype=d), (torch.rand(6)>0.3).double(), True])
run("token_logprobs", [torch.randn(3,5,dtype=d), torch.randint(0,5,(3,))])
run("gae", [torch.randn(4,dtype=d), torch.randn(4,dtype=d), torch.tensor(0.5,dtype=d), 0.99, 0.95])
run("cumulative_max", [torch.randn(6,dtype=d)])
run("median_reward", [torch.randn(7,dtype=d)])

# empty masks
z = torch.zeros(8)
run("masked_mean", [torch.randn(8), z]); run("masked_sum", [torch.randn(8), z])
run("masked_max", [torch.randn(8), z]); run("whiten", [torch.randn(8), z, True])
run("advantage_mean_std", [torch.randn(8), z])
run("entropy", [torch.randn(4,6), torch.zeros(4)])
run("normalized_entropy", [torch.randn(4,6), torch.zeros(4)])
run("sequence_logprob", [torch.randn(3,5,7), torch.randint(0,7,(3,5)), torch.zeros(3,5), True])

# single element
run("cumulative_max", [torch.randn(1)])
run("median_reward", [torch.randn(1)])
run("normalize", [torch.randn(1), 1e-8])
run("mode_label", [torch.randint(0,3,(1,))])

# ties in argmax / mode / top_k
run("argmax_tokens", [torch.tensor([[1.0,3.0,3.0,2.0]])])
run("mode_label", [torch.tensor([0,1,1,2,2])])  # tie between 1 and 2
run("mode_label", [torch.tensor([3,3,1,1,0])])
run("top_k_mask", [torch.tensor([[1.0,2.0,2.0,3.0,3.0]]), 2])

# boundary bucketize
run("bucketize_reward", [torch.tensor([-1.0,0.0,1.0,0.5,-0.5]), torch.tensor([-0.5,0.0,0.5])])

# top_p boundaries
pr = torch.tensor([[0.5,0.3,0.15,0.05]])
for p in [0.0,0.5,0.8,0.5000001,1.0]:
    run("top_p_mask", [pr, p])

# clip_fraction boundary
lp = torch.tensor([0.0, 0.1, -0.1]); olp=torch.zeros(3)
run("clip_fraction", [lp, olp, 0.1])

# quantile edge q=0,1
run("quantile_lower", [torch.randn(9), 0.0]); run("quantile_lower", [torch.randn(9), 1.0])

# discounted / lambda various T
run("discounted_returns", [torch.randn(1), 0.9])
run("lambda_returns", [torch.randn(1), torch.randn(1), torch.tensor(0.3), 0.9, 0.8])

# composites: group_size=1, T=1, beta=0
V=5
run("grpo_objective", [torch.randn(4,1,V),torch.randn(4,1,V),torch.randn(4,1,V),
                       torch.randint(0,V,(4,1)),torch.ones(4,1),torch.randn(4),1,0.0,0.2,0.2,True,"k3"])
run("grpo_objective", [torch.randn(4,3,V),torch.randn(4,3,V),torch.randn(4,3,V),
                       torch.randint(0,V,(4,3)),(torch.rand(4,3)>0.3).float(),torch.randn(4),2,0.1,0.2,0.2,False,"k2"])
run("rloo_objective", [torch.randn(6,3,V),torch.randn(6,3,V),torch.randint(0,V,(6,3)),
                       (torch.rand(6,3)>0.3).float(),torch.randn(6),3,0.2,0.2])
run("reverse_kl_objective", [torch.randn(4,3,V),torch.randn(4,3,V),torch.randint(0,V,(4,3)),
                             (torch.rand(4,3)>0.3).float(),torch.randn(4,3),0.2])
run("ppo_objective", [torch.randn(6),torch.randn(6),torch.randn(6),torch.randn(6),
                      torch.randn(6),torch.tensor(0.2),0.99,1.0,0.2,0.2,0.2,0.5])
run("dpo_sequence_loss", [torch.randn(3,4,V),torch.randn(3,4,V),torch.randn(3,4,V),torch.randn(3,4,V),
                          torch.randint(0,V,(3,4)),torch.randint(0,V,(3,4)),
                          torch.ones(3,4),torch.ones(3,4),0.1,0.2])

print("EDGE FAILURES:", len(fails))
for f in fails: print("  ", f)