File size: 654 Bytes
5ccb4fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | # Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import jax
def snis_mode_baseline(
total,
candidate_log_p,
sampled_log_p,
):
if total.shape != candidate_log_p.shape:
raise ValueError("baseline energy and log-density shapes differ")
if sampled_log_p.shape != total.shape:
raise ValueError("sampled log-density must match baseline energy")
weights = jax.nn.softmax(
candidate_log_p.astype(total.real.dtype)
- sampled_log_p.astype(total.real.dtype),
axis=-1,
)
return weights
__all__ = ["snis_mode_baseline"]
|