| from functools import partial |
|
|
| import jax |
| import jax.numpy as jnp |
|
|
|
|
| @partial(jax.jit, static_argnums=(2,)) |
| def encode_text( |
| input_ids, attention_mask, encoder_apply_fn, encoder_params, |
| latent_mean, latent_std, |
| ): |
| """Encoder pass from text to latent with normalization.""" |
| latents = encoder_apply_fn( |
| {"params": encoder_params}, input_ids=input_ids, attention_mask=attention_mask, |
| deterministic=True, |
| ) |
| return (latents - latent_mean) / latent_std |
|
|
|
|
| def build_self_attn_cond_masks(is_cond, is_valid, xp=jnp): |
| """Build self-attention conditioning masks from cond/valid token flags.""" |
| encoder_attention_mask = ( |
| (is_cond[:, :, None] & is_cond[:, None, :]) | |
| (~is_cond[:, :, None] & is_valid[:, None, :]) |
| ).astype(xp.float32) |
| attention_mask = is_valid.astype(xp.float32) |
| cond_seq_mask = is_cond.astype(xp.float32) |
| return encoder_attention_mask, attention_mask, cond_seq_mask |
|
|
