| # 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"] | |