Fix CI: accept log_rewards kwarg in TBGAFN.get_scores

[Copilot]

test_hypergrid_gafn_smoke was failing with TypeError: TBGAFN.get_scores() got an unexpected keyword argument 'log_rewards' because the method signature in tutorials/examples/train_hypergrid_gafn.py didn't expose that parameter.

Changes

• tutorials/examples/train_hypergrid_gafn.py: Added log_rewards: torch.Tensor | None = None to TBGAFN.get_scores(). Falls back to trajectories.log_rewards when not supplied, preserving existing behavior.

def get_scores(
    self,
    trajectories: Trajectories,
    recalculate_all_logprobs: bool = True,
    log_rewards: torch.Tensor | None = None,
) -> torch.Tensor:
    ...
    if log_rewards is None:
        log_rewards = trajectories.log_rewards
    assert log_rewards is not None

Original prompt

Fix CI failure in PR #503: tutorials/examples/test_scripts.py::test_hypergrid_gafn_smoke fails with TypeError: TBGAFN.get_scores() got an unexpected keyword argument 'log_rewards' (Actions run 23931526330, job 69799350648).

Update the caller-side implementation of TBGAFN.get_scores in tutorials/examples/train_hypergrid_gafn.py to accept a log_rewards keyword argument with default None, and if None fall back to trajectories.log_rewards. Preserve existing behavior and ensure type hints/imports are correct.

Repository: GFNOrg/torchgfn
Base branch: custom-log-rewards
Reference commit for context: 9630594

Acceptance:

• pytest -x passes, specifically the GAFN smoke test.
• No behavior change when log_rewards is not passed.
• Lint/type-check friendly signature.

The following is the prior conversation context from the user's chat exploration (may be truncated):

User: Please find a solution for this failing job. Use the logs, job definition, and any referenced files where the failure occurred. Keep your response focused on the solution and include code suggestions when appropriate.

Assistant: The failing CI job is caused by an API mismatch:

• Failing test: tutorials/examples/test_scripts.py::test_hypergrid_gafn_smoke (lines 569–580 in tutorials/examples/test_scripts.py at ref 9630594bd43b47e5008ba704da9d9c1d294971e0)
• Error from logs (ref 9630594bd43b47e5008ba704da9d9c1d294971e0):
  TypeError: TBGAFN.get_scores() got an unexpected keyword argument 'log_rewards'

That means some caller (very likely the Hypergrid GAFN training example invoked by train_hypergrid_gafn_main) is calling:

algo.get_scores(..., log_rewards=...)

but the implementation of TBGAFN.get_scores() no longer accepts log_rewards (or never did).

Fix (recommended): make TBGAFN.get_scores backward-compatible

Update TBGAFN.get_scores() to accept log_rewards as an optional alias, so older tutorial/example code still runs.

In torchgfn/algorithms/tb_gafn.py (or wherever class TBGAFN is defined), change the signature to include log_rewards=None and map it to whatever the method currently expects (often rewards, log_r, etc.). Example patch:

diff --git a/torchgfn/algorithms/tb_gafn.py b/torchgfn/algorithms/tb_gafn.py
@@
-    def get_scores(self, trajectories, rewards=None, **kwargs):
+    def get_scores(self, trajectories, rewards=None, log_rewards=None, **kwargs):
+        # Backward-compatibility: older callers pass log_rewards
+        if log_rewards is not None and rewards is None:
+            rewards = log_rewards
         ...

Notes:

• If your current implementation expects non-log rewards, then instead do rewards = log_rewards.exp() (or rename the internal variable appropriately). The key is: accept the keyword and convert it as needed.
• Keeping **kwargs is fine, but you still must explicitly list log_rewards if you want to accept it without error.

Alternative fix: update the caller (tutorial script)

If you prefer to enforce the new API, change the call site in the Hypergrid GAFN training example (the function imported as train_hypergrid_gafn_main in tutorials/examples/test_scripts.py line 45). Replace log_rewards=... with the new expected argument name.

However, given this is a tutorial/example smoke test, the compatibility fix in TBGAFN.get_scores() is usually the safest and least disruptive: it preserves older scripts and avoids breaking downstream user code.

Why this will unblock CI

The workflow runs pytest -x (see .github/workflows/ci.yml at ref 9630594bd43b47e5008ba704da9d9c1d294971e0, lines 30–36), and the run stopped at the first failure. Fixing this one signature mismatch should let the suite continue (and it was otherwise very healthy: 1641 passed, 700 skipped).

User: can you update the TBGAFN.get_scores() caller with default value of None?

Assistant: Update the caller-side TBGAFN.get_scores() signature to accept log_rewards with a default of None, so any existing code path that passes log_rewards=... won’t crash, while keeping current behavior unchanged when it’s not provided.

The CI failure was:

• tutorials/examples/test_scripts.py::test_hypergrid_gafn_smoke
• TypeError: TBGAFN.get_scores() got an unexpected keyword argument 'log_rewards' (job logs at ref 9630594bd43b47e5008ba704da9d9c1d294971e0)

Concrete code change

In tutorials/examples/train_hypergrid_gafn.py, modify TBGAFN.get_scores to accept log_rewards: torch.Tensor | None = None and fall back to trajectories.log_rewards when log_rewards isn’t provided:

def get_scores(
    self,
    trajectories: Trajectories,
    recalculate_all_logprobs: bool = True,
    log_rewards: torch.Tenso...

</details>



<!-- START COPILOT CODING AGENT SUFFIX -->

*This pull request was created from Copilot chat.*
>

[josephdviviano]

lgtm!

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (custom-log-rewards@9630594). Learn more about missing BASE report.

Additional details and impacted files

@@                  Coverage Diff                  @@
##             custom-log-rewards     #504   +/-   ##
=====================================================
  Coverage                      ?   72.46%           
=====================================================
  Files                         ?       55           
  Lines                         ?     8537           
  Branches                      ?     1098           
=====================================================
  Hits                          ?     6186           
  Misses                        ?     1961           
  Partials                      ?      390           

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.