torchcell.losses.distributional.make_dist_head

make_dist_head(dist: str, num_quantiles: int = 19, num_features: int | None = None, rank: int = 32, num_samples: int = 128) → DistHead[source]

Build a DistHead from a config dist value.

Parameters:

  • dist – One of point / crps / quantile / laplace_crps / nll / energy (the keys of DIST_TO_MODE).

  • num_quantiles – quantile only – size of the tau grid.

  • num_features – energy only – F; REQUIRED there to allocate V [F, rank].

  • rank – energy only – k (0 = the diagonal-only ablation).

  • num_samples – energy only – m predictive samples per score.

Returns:

The configured head.