File size: 3,292 Bytes
d766458 | 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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | import jax
import jax.numpy as jnp
import haiku as hk
import numpy as np
from .utils import cat_neighbors_nodes, get_ar_mask
class mpnn_sample:
def sample(self, I):
"""
I = {
[[required]]
'X' = (L,4,3)
'mask' = (L,)
'residue_index' = (L,)
'chain_idx' = (L,)
'decoding_order' = (L,)
[[optional]]
'ar_mask' = (L,L)
'bias' = (L,21)
'temperature' = 1.0
}
"""
key = hk.next_rng_key()
L = I["X"].shape[0]
temperature = I.get("temperature",1.0)
# prepare node and edge embeddings
E, E_idx = self.features(I)
h_V = jnp.zeros((E.shape[0], E.shape[-1]))
h_E = self.W_e(E)
##############
# encoder
##############
mask_attend = jnp.take_along_axis(I["mask"][:,None] * I["mask"][None,:], E_idx, 1)
for layer in self.encoder_layers:
h_V, h_E = layer(h_V, h_E, E_idx, I["mask"], mask_attend)
# get autoregressive mask
ar_mask = I.get("ar_mask",get_ar_mask(I["decoding_order"]))
mask_attend = jnp.take_along_axis(ar_mask, E_idx, 1)
mask_1D = I["mask"][:,None]
mask_bw = mask_1D * mask_attend
mask_fw = mask_1D * (1 - mask_attend)
h_EX_encoder = cat_neighbors_nodes(jnp.zeros_like(h_V), h_E, E_idx)
h_EXV_encoder = cat_neighbors_nodes(h_V, h_EX_encoder, E_idx)
h_EXV_encoder = mask_fw[...,None] * h_EXV_encoder
def fwd(x, t, key):
h_EXV_encoder_t = h_EXV_encoder[t]
E_idx_t = E_idx[t]
mask_t = I["mask"][t]
mask_bw_t = mask_bw[t]
h_ES_t = cat_neighbors_nodes(x["h_S"], h_E[t], E_idx_t)
##############
# decoder
##############
for l,layer in enumerate(self.decoder_layers):
h_V = x["h_V"][l]
h_ESV_decoder_t = cat_neighbors_nodes(h_V, h_ES_t, E_idx_t)
h_ESV_t = mask_bw_t[...,None] * h_ESV_decoder_t + h_EXV_encoder_t
h_V_t = layer(h_V[t], h_ESV_t, mask_V=mask_t)
# update
x["h_V"] = x["h_V"].at[l+1,t].set(h_V_t)
logits_t = self.W_out(h_V_t)
x["logits"] = x["logits"].at[t].set(logits_t)
##############
# sample
##############
# add bias
if "bias" in I: logits_t += I["bias"][t]
# sample character
logits_t = logits_t/temperature + jax.random.gumbel(key, logits_t.shape)
# tie positions
logits_t = logits_t.mean(0, keepdims=True)
S_t = jax.nn.one_hot(logits_t[...,:20].argmax(-1), 21)
# update
x["h_S"] = x["h_S"].at[t].set(self.W_s(S_t))
x["S"] = x["S"].at[t].set(S_t)
return x, None
# initial values
X = {"h_S": jnp.zeros_like(h_V),
"h_V": jnp.array([h_V] + [jnp.zeros_like(h_V)] * len(self.decoder_layers)),
"S": jnp.zeros((L,21)),
"logits": jnp.zeros((L,21))}
# scan over decoding order
t = I["decoding_order"]
if t.ndim == 1: t = t[:,None]
XS = {"t":t, "key":jax.random.split(key,t.shape[0])}
X = hk.scan(lambda x, xs: fwd(x, xs["t"], xs["key"]), X, XS)[0]
return {"S":X["S"], "logits":X["logits"], "decoding_order":t} |