How to use the jax.numpy.zeros function in jax
To help you get started, we’ve selected a few jax examples, based on popular ways it is used in public projects.
Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
JuliusKunze / jaxnet / tests / test_jaxnet.py View on Github 
def test_external_submodule2():
layer = Dense(2, zeros, zeros)
@parametrized
def net(inputs):
return layer(inputs)
inputs = np.zeros((1, 2))
params = net.init_params(PRNGKey(0), inputs)
assert_params_equal(((np.zeros((2, 2)), np.zeros(2)),), params)
out = net.apply(params, inputs)
assert np.array_equal(np.zeros((1, 2)), out)
out_ = jit(net.apply)(params, inputs)
assert np.array_equal(out, out_)def test_external_submodule2():
layer = Dense(2, zeros, zeros)
@parametrized
def net(inputs):
return layer(inputs)
inputs = np.zeros((1, 2))
params = net.init_params(PRNGKey(0), inputs)
assert_params_equal(((np.zeros((2, 2)), np.zeros(2)),), params)
out = net.apply(params, inputs)
assert np.array_equal(np.zeros((1, 2)), out)
out_ = jit(net.apply)(params, inputs)
assert np.array_equal(out, out_)def test_no_params():
@parametrized
def double(inputs):
return 2 * inputs
inputs = np.zeros((1, 3))
params = double.init_params(PRNGKey(0), inputs)
assert_params_equal((), params)
out = double.apply(params, inputs)
assert np.array_equal(np.zeros((1, 3)), out)
out_ = jit(double.apply)(params, inputs)
assert np.array_equal(out, out_)def test_external_submodule2():
layer = Dense(2, zeros, zeros)
@parametrized
def net(inputs):
return layer(inputs)
inputs = np.zeros((1, 2))
params = net.init_params(PRNGKey(0), inputs)
assert_params_equal(((np.zeros((2, 2)), np.zeros(2)),), params)
out = net.apply(params, inputs)
assert np.array_equal(np.zeros((1, 2)), out)
out_ = jit(net.apply)(params, inputs)
assert np.array_equal(out, out_)def test_ocr_rnn():
length = 5
carry_size = 3
class_count = 4
inputs = np.zeros((1, length, 4))
def rnn(): return Rnn(*GRUCell(carry_size, zeros))
net = Sequential(
rnn(),
rnn(),
rnn(),
lambda x: np.reshape(x, (-1, carry_size)), # -> same weights for all time steps
Dense(class_count, zeros, zeros),
softmax,
lambda x: np.reshape(x, (-1, length, class_count)))
params = net.init_parameters(inputs, key=PRNGKey(0))
assert len(params) == 4
cell = params.rnn0.gru_celldef test_external_param_sharing():
layer = Dense(2, zeros, zeros)
shared_net = Sequential(layer, layer)
inputs = np.zeros((1, 2))
params = shared_net.init_params(PRNGKey(0), inputs)
assert_params_equal(((np.zeros((2, 2)), np.zeros(2)),), params)def test_parameter_with_multiple_arrays_submodule():
@parametrized
def wrapper():
return Parameter(lambda _: (np.zeros(()), np.zeros(())))()
params = wrapper.init_parameters(key=PRNGKey(0))
a, b = params.parameter
assert np.zeros(()) == a
assert np.zeros(()) == b
out = wrapper.apply(params)
assert params.parameter == outdef init_fn(_):
return ScaleByScheduleState(count=jnp.zeros([]))def compute_pairwise_mean_variance(X, Y, dim1, dim2, eta1, eta2, c, kappa, omega):
probe = np.zeros((4, X.shape[1]))
probe = jax.ops.index_update(probe, jax.ops.index[:, dim1], np.array([1.0, 1.0, -1.0, -1.0]))
probe = jax.ops.index_update(probe, jax.ops.index[:, dim2], np.array([1.0, -1.0, 1.0, -1.0]))
vec = np.array([0.25, -0.25, -0.25, 0.25])
return compute_coefficient_mean_variance(X, Y, probe, vec, eta1, eta2, c, kappa, omega)def model(X, Y):
# set uninformative log-normal priors on our three kernel hyperparameters
var = numpyro.sample("kernel_var", dist.LogNormal(0.0, 10.0))
noise = numpyro.sample("kernel_noise", dist.LogNormal(0.0, 10.0))
length = numpyro.sample("kernel_length", dist.LogNormal(0.0, 10.0))
# compute kernel
k = kernel(X, X, var, length, noise)
# sample Y according to the standard gaussian process formula
numpyro.sample("Y", dist.MultivariateNormal(loc=np.zeros(X.shape[0]), covariance_matrix=k),
obs=Y)jax
Differentiate, compile, and transform Numpy code.
