Step functions fixes, add in_context_pvi
#223
Merged
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR:
kl_divergence
step function with a more general utility functioncontrast_logit_fn
andcontrast_logit_diff_fn
, equivalent tocontrast_prob_fn
andcontrast_prob_diff_fn
but using pre-softmax logits instead of probabilities.logprobs
boolean parameter toprobability_fn
,contrast_prob_fn
,contrast_prob_diff_fn
andmc_dropout_prob_avg_fn
to allow for usingF.log_softmax
log-probabilities instead ofF.softmax
probabilities (default behavior)in_context_pvi
step function for calculating in-context pointwise V-usable information (Lu et al. 2023), using @boblus implementation from boblus/in-context-pvi