jcm/losses.py [751:786]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        loss = jnp.nansum(losses * batch["mask"] / jnp.sum(batch["mask"]))
        quarter_masks = get_quarter_masks(
            t,
            np.linspace(sde.t_min ** (1 / sde.rho), sde.t_max ** (1 / sde.rho), 5)
            ** sde.rho,
        )
        loss_q1 = jnp.nansum(
            losses
            * quarter_masks[0]
            * batch["mask"]
            / jnp.sum(quarter_masks[0] * batch["mask"])
        )
        loss_q2 = jnp.nansum(
            losses
            * quarter_masks[1]
            * batch["mask"]
            / jnp.sum(quarter_masks[1] * batch["mask"])
        )
        loss_q3 = jnp.nansum(
            losses
            * quarter_masks[2]
            * batch["mask"]
            / jnp.sum(quarter_masks[2] * batch["mask"])
        )
        loss_q4 = jnp.nansum(
            losses
            * quarter_masks[3]
            * batch["mask"]
            / jnp.sum(quarter_masks[3] * batch["mask"])
        )

        log_stats = {
            "loss_q1": loss_q1,
            "loss_q2": loss_q2,
            "loss_q3": loss_q3,
            "loss_q4": loss_q4,
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



jcm/losses.py [848:883]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        loss = jnp.nansum(losses * batch["mask"] / jnp.sum(batch["mask"]))
        quarter_masks = get_quarter_masks(
            t,
            np.linspace(sde.t_min ** (1 / sde.rho), sde.t_max ** (1 / sde.rho), 5)
            ** sde.rho,
        )
        loss_q1 = jnp.nansum(
            losses
            * quarter_masks[0]
            * batch["mask"]
            / jnp.sum(quarter_masks[0] * batch["mask"])
        )
        loss_q2 = jnp.nansum(
            losses
            * quarter_masks[1]
            * batch["mask"]
            / jnp.sum(quarter_masks[1] * batch["mask"])
        )
        loss_q3 = jnp.nansum(
            losses
            * quarter_masks[2]
            * batch["mask"]
            / jnp.sum(quarter_masks[2] * batch["mask"])
        )
        loss_q4 = jnp.nansum(
            losses
            * quarter_masks[3]
            * batch["mask"]
            / jnp.sum(quarter_masks[3] * batch["mask"])
        )

        log_stats = {
            "loss_q1": loss_q1,
            "loss_q2": loss_q2,
            "loss_q3": loss_q3,
            "loss_q4": loss_q4,
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -