in jcm/models/normalization.py [0:0]
def __call__(self, x, y):
means = jnp.mean(x, axis=(1, 2))
m = jnp.mean(means, axis=-1, keepdims=True)
v = jnp.var(means, axis=-1, keepdims=True)
means_plus = (means - m) / jnp.sqrt(v + 1e-5)
h = (x - means[:, None, None, :]) / jnp.sqrt(
jnp.var(x, axis=(1, 2), keepdims=True) + 1e-5
)
normal_init = init.normal(0.02)
zero_init = init.zeros
if self.bias:
def init_embed(key, shape, dtype=jnp.float32):
feature_size = shape[1] // 3
normal = (
normal_init(key, (shape[0], 2 * feature_size), dtype=dtype) + 1.0
)
zero = zero_init(key, (shape[0], feature_size), dtype=dtype)
return jnp.concatenate([normal, zero], axis=-1)
embed = nn.Embed(
num_embeddings=self.num_classes,
features=x.shape[-1] * 3,
embedding_init=init_embed,
)
else:
def init_embed(key, shape, dtype=jnp.float32):
return normal_init(key, shape, dtype=dtype) + 1.0
embed = nn.Embed(
num_embeddings=self.num_classes,
features=x.shape[-1] * 2,
embedding_init=init_embed,
)
if self.bias:
gamma, alpha, beta = jnp.split(embed(y), 3, axis=-1)
h = h + means_plus[:, None, None, :] * alpha[:, None, None, :]
out = gamma[:, None, None, :] * h + beta[:, None, None, :]
else:
gamma, alpha = jnp.split(embed(y), 2, axis=-1)
h = h + means_plus[:, None, None, :] * alpha[:, None, None, :]
out = gamma[:, None, None, :] * h
return out