Spaces:
Sleeping
Sleeping
File size: 1,908 Bytes
3dce43d d516303 3dce43d |
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 |
from functools import partial
from typing import Callable
import jax
import jax.numpy as jnp
import tensorflow_probability.substrates.jax as tfp
from flax import linen as nn
tfd = tfp.distributions
def split_tree(a, rng_key):
treedef = jax.tree_util.tree_structure(a)
num_vars = len(jax.tree_util.tree_leaves(a))
all_keys = jax.random.split(rng_key, num=(num_vars + 1))
return jax.tree_util.tree_unflatten(treedef, all_keys[1:])
def sample_fn(rng, vi_params: nn.FrozenDict):
rng = split_tree(vi_params["mean"], rng)
params = jax.tree_map(
lambda m, ls, k: tfd.Normal(loc=m, scale=jnp.exp(ls)).sample(seed=k),
vi_params["mean"],
vi_params["log_scale"],
rng,
) # type: nn.FrozenDict
return params
def get_apply_fn(model: nn.Module):
"""Returns the model forward function"""
@jax.jit
@partial(jax.vmap, in_axes=(None, None, 0))
def apply_fn(vi_params, inputs, rng):
params = sample_fn(rng, vi_params)
outputs = model.apply({"params": params}, inputs)
return outputs
@jax.jit
def apply_map_fn(params, inputs, rng):
outputs = model.apply({"params": params["mean"]}, inputs)
return outputs[None, ...]
return apply_fn, apply_map_fn
class MLP(nn.Module):
n_features: int = 512
n_layers: int = 3
n_classes: int = 10
n_features_mult: int = 1
bias_init: Callable = nn.initializers.zeros_init()
act: Callable = nn.relu
dtype: str = "float32"
@nn.compact
def __call__(self, x):
dense = partial(
nn.Dense,
dtype=self.dtype,
bias_init=self.bias_init,
)
x = jnp.reshape(x, (x.shape[0], -1))
for _ in range(self.n_layers):
x = dense(int(self.n_features * self.n_features_mult))(x)
x = nn.relu(x)
x = dense(self.n_classes)(x)
return x
|