Support callable and per-state policy temperatures

[alstn12088]

Summary

This PR extends discrete policy sampling so callers can provide a state-dependent sampling temperature instead of only a scalar temperature.

Concretely, it:

• lets PolicyMixin and RecurrentPolicyMixin accept policy_kwargs["temperature"] as a callable
• calls that temperature function after estimator outputs are available, with (states_active, estimator_outputs)
• allows LogitBasedEstimator / DiscretePolicyEstimator temperature to be either a scalar float or a tensor
• broadcasts tensor temperatures across the action/logit dimension before dividing logits
• preserves the existing scalar-temperature path unchanged

Motivation

The downstream gfn-communities novelty-guidance work needs local, per-state adaptive temperature during trajectory sampling. The sampler already threads temperature through policy_kwargs; this change makes that hook expressive enough for per-step, state-conditioned temperature without forcing downstream code to fork the sampler or estimator stack.

Behavioral Compatibility

• Existing scalar temperature=1.0 and temperature=<float> behavior is unchanged.
• Callable temperature is opt-in. If temperature is not callable, the old path is used.
• Tensor temperature is opt-in. Scalar comparisons/assertions are preserved for scalar values, while tensor values assert all entries are positive.
• The tensor path only changes logits by dividing by a broadcasted positive tensor.

Testing

Local smoke run from the downstream environment:

PYTHONPATH=src /home/mila/m/minsu.kim/.conda/envs/gfn-communities/bin/python - <<'PY'
import torch
from gfn.estimators import LogitBasedEstimator
logits = torch.tensor([[1.0, 2.0], [3.0, 5.0]])
masks = torch.ones_like(logits, dtype=torch.bool)
out = LogitBasedEstimator._prepare_logits(
    logits, masks, sf_index=None, sf_bias=0.0, temperature=torch.tensor([1.0, 2.0])
)
assert torch.allclose(out[0], logits[0])
assert torch.allclose(out[1], logits[1] / 2.0)
scalar = LogitBasedEstimator._prepare_logits(
    logits, masks, sf_index=None, sf_bias=0.0, temperature=2.0
)
assert torch.allclose(scalar, logits / 2.0)
print('tensor/scalar temperature smoke passed')
PY

Output:

tensor/scalar temperature smoke passed

Downstream Dependency

This PR is required by the gfn-communities adaptive-temperature novelty guidance implementation. The downstream core PR intentionally keeps the submodule pointer unchanged until this lands in GFNOrg/torchgfn; after merge, the submodule pointer should be updated in a follow-up.

[codecov]

Codecov Report

❌ Patch coverage is 16.66667% with 15 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.43%. Comparing base (b1571a0) to head (2162a02).
⚠️ Report is 17 commits behind head on master.

Files with missing lines	Patch %	Lines
src/gfn/estimators.py	16.66%	10 Missing and 5 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #520      +/-   ##
==========================================
+ Coverage   73.91%   74.43%   +0.52%     
==========================================
  Files          59       59              
  Lines        9730     9906     +176     
  Branches     1340     1367      +27     
==========================================
+ Hits         7192     7374     +182     
+ Misses       2095     2073      -22     
- Partials      443      459      +16     

Files with missing lines	Coverage Δ
src/gfn/estimators.py	85.91% <16.66%> (-3.09%)	⬇️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[josephdviviano]

great work - thanks. forcing the merge to bypass the weird claude review issue