in jcm/models/utils.py [0:0]
def get_score_fn(sde, model, params, states, train=False, return_state=False):
"""Wraps `score_fn` so that the model output corresponds to a real time-dependent score function.
Args:
sde: An `sde_lib.SDE` object that represents the forward SDE.
model: A `flax.linen.Module` object that represents the architecture of the score-based model.
params: A dictionary that contains all trainable parameters.
states: A dictionary that contains all other mutable parameters.
train: `True` for training and `False` for evaluation.
return_state: If `True`, return the new mutable states alongside the model output.
Returns:
A score function.
"""
model_fn = get_model_fn(model, params, states, train=train)
if isinstance(sde, sde_lib.VPSDE) or isinstance(sde, sde_lib.subVPSDE):
def score_fn(x, t, rng=None):
# Scale neural network output by standard deviation and flip sign
# For VP-trained models, t=0 corresponds to the lowest noise level
# The maximum value of time embedding is assumed to 999 for
# continuously-trained models.
cond_t = t * 999
model, state = model_fn(x, cond_t, rng)
std = sde.marginal_prob(jnp.zeros_like(x), t)[1]
score = batch_mul(-model, 1.0 / std)
if return_state:
return score, state
else:
return score
elif isinstance(sde, sde_lib.VESDE):
def score_fn(x, t, rng=None):
x = 2 * x - 1.0 # assuming x is in [0, 1]
std = sde.marginal_prob(jnp.zeros_like(x), t)[1]
score, state = model_fn(x, jnp.log(std), rng)
score = batch_mul(score, 1.0 / std)
if return_state:
return score, state
else:
return score
elif isinstance(sde, sde_lib.KVESDE):
denoiser_fn = get_denoiser_fn(
sde, model, params, states, train=train, return_state=True
)
def score_fn(x, t, rng=None):
denoiser, state = denoiser_fn(x, t, rng)
score = batch_mul(denoiser - x, 1 / t**2)
if return_state:
return score, state
else:
return score
else:
raise NotImplementedError(
f"SDE class {sde.__class__.__name__} not yet supported."
)
return score_fn