def conv_or_conv_transpose(inputs):
V = parameter(filter_shape + (inputs.shape[-1], out_chan), normal(.05), 'V')
example_out = apply(inputs, V=V, g=np.ones(out_chan), b=np.zeros(out_chan))
# TODO remove need for `.aval.val` when capturing variables in initializer function:
g = Parameter(lambda key: init_scale /
np.sqrt(np.var(example_out.aval.val, (0, 1, 2)) + 1e-10), 'g')()
b = Parameter(lambda key: np.mean(example_out.aval.val, (0, 1, 2)) * g.aval.val, 'b')()
return apply(inputs, V, b, g)
def cZ(self, uv):
u, v = uv
return self.sign * jnp.sqrt(1.0 - u**2 - v**2)
def _Z(self, uv):
w, B = uv
a, b, c = self.abc
top = (a**2 * jnp.sin(w)**2) + (b**2 * jnp.cos(w)**2) - c**2
return c * jnp.sin(B) * (jnp.sqrt(top) / jnp.sqrt(a**2 - c**2))
def _X(self, uv):
w, B = uv
a, b, c = self.abc
top = a**2 - (b**2 * jnp.sin(B)**2) - (c**2 * jnp.cos(B)**2)
return a * jnp.cos(w) * (jnp.sqrt(top) / jnp.sqrt(a**2 - c**2))
def vector_magnitudes(vecs):
""" vectorized, differentiable function for vector
magnitudes
"""
return jnp.sqrt(jnp.einsum("...i,...i", vecs, vecs))
def _l2_normalize(arr, axis):
return arr / np.sqrt(np.sum(arr**2, axis=axis, keepdims=True))