jejjohnson / jejjohnson/pipekit
pipekit-train: KL sample-based fallback (v0.1.1)
- Dominant language
- Python
- Stars
- 0
- Forks
- 0
- PR merge metrics
- No merged PRs in 30d
Description
## Problem / Request
The `KL` loss currently supports only the **analytic** path (`predicted.kl_divergence(target)`). The design committed to a sample-based fallback `-(target.log_prob(samples) - predicted.log_prob(samples)).mean()` with samples drawn from `predicted`. Deferred from v0.1 because it needs the per-batch RNG; landed under v0.1.x.
## Motivation
- Many distributions don't expose `.kl_divergence` (custom densities, mixture distributions, normalising flows). The sample-based fallback unblocks them.
- Small, contained change: just needs an RNG threaded through `Loss.__call__`.
## Proposed API
Today's `Loss.__call__(predicted, target)` extends to `Loss.__call__(predicted, target, key=None)`:
```python
class KL(Operator):
n_samples: int = 32 # for the sample-based fallback
def _apply(self, predicted, target, key=None):
if hasattr(predicted, "kl_divergence"):
return predicted.kl_divergence(target).mean(), {"kl": ...}
if key is None:
raise TypeError(
"Sample-based KL requires a PRNG key. The Equinox "
"adapter supplies one via task.loss_fn; for ad-hoc "
"use, pass key=jax.random.key(seed) explicitly."
)
samples = predicted.sample(seed=key, sample_shape=(self.n_samples,))
log_q = predicted.log_prob(samples)
log_p = target.log_prob(samples)
return (log_q - log_p).mean(), {"kl": ...}
```
`_SynthesisedTask.loss_fn` in the Equinox adapter already has `key`; thread it through.
## Implementation Steps
- [ ] Extend `Loss.__call__` signature to accept optional `key` (backwards compatible — default `None`).
- [ ] `KL._apply` implements the sample-based fallback.
- [ ] `_SynthesisedTask.loss_fn` passes `key` to the loss call.
- [ ] Test: KL between two simple distributions (e.g. mocked Normal + Normal) via the analytic path; same pair via the sample-based path; assert the two agree within tolerance.
## Definition of Done
- [ ] Sample-based KL works under the Equinox adapter.
- [ ] Existing analytic path unchanged.
- [ ] Docstring updated; `KL` no longer raises `TypeError` for non-analytic distributions.
## Relationships
- Parent: pipekit-train v0.1.x (small follow-up)
Contributor guide
Assessment
This issue has not been assessed yet.