Python jax.numpy.zeros() Examples
The following are 30
code examples of jax.numpy.zeros().
You can vote up the ones you like or vote down the ones you don't like,
and go to the original project or source file by following the links above each example.
You may also want to check out all available functions/classes of the module
jax.numpy
, or try the search function
.
![](https://www.programcreek.com/common/static/images/search.png)
Example #1
Source File: hmm.py From numpyro with Apache License 2.0 | 7 votes |
def forward_log_prob(init_log_prob, words, transition_log_prob, emission_log_prob, unroll_loop=False): # Note: The following naive implementation will make it very slow to compile # and do inference. So we use lax.scan instead. # # >>> log_prob = init_log_prob # >>> for word in words: # ... log_prob = forward_one_step(log_prob, word, transition_log_prob, emission_log_prob) def scan_fn(log_prob, word): return forward_one_step(log_prob, word, transition_log_prob, emission_log_prob), jnp.zeros((0,)) if unroll_loop: log_prob = init_log_prob for word in words: log_prob = forward_one_step(log_prob, word, transition_log_prob, emission_log_prob) else: log_prob, _ = lax.scan(scan_fn, init_log_prob, words) return log_prob
Example #2
Source File: proportion_test.py From numpyro with Apache License 2.0 | 6 votes |
def make_dataset(rng_key) -> Tuple[jnp.ndarray, jnp.ndarray]: """ Make simulated dataset where potential customers who get a sales calls have ~2% higher chance of making another purchase. """ key1, key2, key3 = random.split(rng_key, 3) num_calls = 51342 num_no_calls = 48658 made_purchase_got_called = dist.Bernoulli(0.084).sample(key1, sample_shape=(num_calls,)) made_purchase_no_calls = dist.Bernoulli(0.061).sample(key2, sample_shape=(num_no_calls,)) made_purchase = jnp.concatenate([made_purchase_got_called, made_purchase_no_calls]) is_female = dist.Bernoulli(0.5).sample(key3, sample_shape=(num_calls + num_no_calls,)) got_called = jnp.concatenate([jnp.ones(num_calls), jnp.zeros(num_no_calls)]) design_matrix = jnp.hstack([jnp.ones((num_no_calls + num_calls, 1)), got_called.reshape(-1, 1), is_female.reshape(-1, 1)]) return design_matrix, made_purchase
Example #3
Source File: test_core.py From jaxnet with Apache License 2.0 | 6 votes |
def test_mixed_up_execution_order(): @parametrized def dense(inputs): bias = parameter((2,), zeros, 'bias') kernel = parameter((inputs.shape[-1], 2), zeros, 'kernel') return jnp.dot(inputs, kernel) + bias inputs = jnp.zeros((1, 3)) params = dense.init_parameters(inputs, key=PRNGKey(0)) assert (2,) == params.bias.shape assert (3, 2) == params.kernel.shape out = dense.apply(params, inputs) assert jnp.array_equal(jnp.zeros((1, 2)), out) out_ = dense.apply(params, inputs, jit=True) assert jnp.array_equal(out, out_)
Example #4
Source File: test_core.py From jaxnet with Apache License 2.0 | 6 votes |
def test_submodule_reuse(): inputs = jnp.zeros((1, 2)) layer = Dense(5) net1 = Sequential(layer, Dense(2)) net2 = Sequential(layer, Dense(3)) layer_params = layer.init_parameters(inputs, key=PRNGKey(0)) net1_params = net1.init_parameters(inputs, key=PRNGKey(1), reuse={layer: layer_params}) net2_params = net2.init_parameters(inputs, key=PRNGKey(2), reuse={layer: layer_params}) out1 = net1.apply(net1_params, inputs) assert out1.shape == (1, 2) out2 = net2.apply(net2_params, inputs) assert out2.shape == (1, 3) assert_dense_parameters_equal(layer_params, net1_params[0]) assert_dense_parameters_equal(layer_params, net2_params[0]) new_layer_params = layer.init_parameters(inputs, key=PRNGKey(3)) combined_params = net1.parameters_from({net1: net1_params, layer: new_layer_params}, inputs) assert_dense_parameters_equal(new_layer_params, combined_params.dense0) assert_dense_parameters_equal(net1_params.dense1, combined_params.dense1)
Example #5
Source File: test_core.py From jaxnet with Apache License 2.0 | 6 votes |
def test_parameters_from(): layer = Dense(2) net = Sequential(layer, relu) inputs = jnp.zeros((1, 3)) layer_params = layer.init_parameters(inputs, key=PRNGKey(0)) params_ = net.parameters_from({layer: layer_params}, inputs) assert_parameters_equal((layer_params,), params_) out = net.apply(params_, inputs) out_ = net.apply_from({layer: layer_params}, inputs) assert jnp.array_equal(out, out_) out_ = net.apply_from({layer: layer_params}, inputs, jit=True) assert jnp.array_equal(out, out_)
Example #6
Source File: test_core.py From jaxnet with Apache License 2.0 | 6 votes |
def test_external_sequential_submodule(): layer = Sequential(Conv(4, (2, 2)), flatten, relu, Dense(3), relu, Dense(2), Sequential(Dense(2), relu)) inputs = jnp.zeros((1, 5, 5, 2)) params = layer.init_parameters(inputs, key=PRNGKey(0)) assert (4,) == params.conv.bias.shape assert (3,) == params.dense0.bias.shape assert (3, 2) == params.dense1.kernel.shape assert (2,) == params.dense1.bias.shape assert (2,) == params.sequential.dense.bias.shape out = layer.apply(params, inputs) assert (1, 2) == out.shape out_ = layer.apply(params, inputs, jit=True) assert jnp.allclose(out, out_)
Example #7
Source File: test_core.py From jaxnet with Apache License 2.0 | 6 votes |
def test_external_submodule2(): layer = Dense(2, zeros, zeros) @parametrized def net(inputs): return layer(inputs) inputs = jnp.zeros((1, 2)) params = net.init_parameters(inputs, key=PRNGKey(0)) assert_parameters_equal(((jnp.zeros((2, 2)), jnp.zeros(2)),), params) out = net.apply(params, inputs) assert jnp.array_equal(jnp.zeros((1, 2)), out) out_ = net.apply(params, inputs, jit=True) assert jnp.array_equal(out, out_)
Example #8
Source File: test_core.py From jaxnet with Apache License 2.0 | 6 votes |
def test_parameters_from_subsubmodule(): subsublayer = Dense(2) sublayer = Sequential(subsublayer, relu) net = Sequential(sublayer, jnp.sum) inputs = jnp.zeros((1, 3)) params = net.init_parameters(inputs, key=PRNGKey(0)) out = net.apply(params, inputs) subsublayer_params = subsublayer.init_parameters(inputs, key=PRNGKey(0)) params_ = net.parameters_from({subsublayer: subsublayer_params}, inputs) assert_dense_parameters_equal(subsublayer_params, params_[0][0]) out_ = net.apply(params_, inputs) assert out.shape == out_.shape out_ = net.apply_from({subsublayer: subsublayer_params}, inputs) assert out.shape == out_.shape out_ = net.apply_from({subsublayer: subsublayer_params}, inputs, jit=True) assert out.shape == out_.shape
Example #9
Source File: test_core.py From jaxnet with Apache License 2.0 | 6 votes |
def test_param_and_submodule_mixed(): @parametrized def linear_map(inputs): kernel = parameter((inputs.shape[-1], 2), zeros, 'kernel') return jnp.dot(inputs, kernel) @parametrized def dense(inputs): return linear_map(inputs) + parameter((2,), zeros, 'bias') inputs = jnp.zeros((1, 3)) params = dense.init_parameters(inputs, key=PRNGKey(0)) assert (2,) == params.bias.shape assert (3, 2) == params.linear_map.kernel.shape out = dense.apply(params, inputs) assert jnp.array_equal(jnp.zeros((1, 2)), out) out_ = dense.apply(params, inputs, jit=True) assert jnp.array_equal(out, out_)
Example #10
Source File: test_core.py From jaxnet with Apache License 2.0 | 6 votes |
def test_submodule_order(): @parametrized def net(): p = Parameter(lambda key: jnp.zeros((1,))) a = p() b = parameter((2,), zeros) c = parameter((3,), zeros) d = parameter((4,), zeros) e = parameter((5,), zeros) f = parameter((6,), zeros) # must not mess up order (decided by first submodule call): k = p() return jnp.concatenate([a, f]) + jnp.concatenate([b, e]) + jnp.concatenate([c, d]) + k params = net.init_parameters(key=PRNGKey(0)) assert jnp.zeros((1,)) == params.parameter0 out = net.apply(params) assert (7,) == out.shape
Example #11
Source File: test_core.py From jaxnet with Apache License 2.0 | 6 votes |
def test_parameters_from_sharing_between_multiple_parents(): a = Dense(2) b = Sequential(a, jnp.sum) @parametrized def net(inputs): return a(inputs), b(inputs) inputs = jnp.zeros((1, 3)) a_params = a.init_parameters(inputs, key=PRNGKey(0)) out = a.apply(a_params, inputs) params = net.parameters_from({a: a_params}, inputs) assert_dense_parameters_equal(a_params, params.dense) assert_parameters_equal((), params.sequential) assert 2 == len(params) out_, _ = net.apply(params, inputs) assert jnp.array_equal(out, out_)
Example #12
Source File: bnn.py From numpyro with Apache License 2.0 | 6 votes |
def model(X, Y, D_H): D_X, D_Y = X.shape[1], 1 # sample first layer (we put unit normal priors on all weights) w1 = numpyro.sample("w1", dist.Normal(jnp.zeros((D_X, D_H)), jnp.ones((D_X, D_H)))) # D_X D_H z1 = nonlin(jnp.matmul(X, w1)) # N D_H <= first layer of activations # sample second layer w2 = numpyro.sample("w2", dist.Normal(jnp.zeros((D_H, D_H)), jnp.ones((D_H, D_H)))) # D_H D_H z2 = nonlin(jnp.matmul(z1, w2)) # N D_H <= second layer of activations # sample final layer of weights and neural network output w3 = numpyro.sample("w3", dist.Normal(jnp.zeros((D_H, D_Y)), jnp.ones((D_H, D_Y)))) # D_H D_Y z3 = jnp.matmul(z2, w3) # N D_Y <= output of the neural network # we put a prior on the observation noise prec_obs = numpyro.sample("prec_obs", dist.Gamma(3.0, 1.0)) sigma_obs = 1.0 / jnp.sqrt(prec_obs) # observe data numpyro.sample("Y", dist.Normal(z3, sigma_obs), obs=Y) # helper function for HMC inference
Example #13
Source File: ucbadmit.py From numpyro with Apache License 2.0 | 6 votes |
def glmm(dept, male, applications, admit=None): v_mu = numpyro.sample('v_mu', dist.Normal(0, jnp.array([4., 1.]))) sigma = numpyro.sample('sigma', dist.HalfNormal(jnp.ones(2))) L_Rho = numpyro.sample('L_Rho', dist.LKJCholesky(2, concentration=2)) scale_tril = sigma[..., jnp.newaxis] * L_Rho # non-centered parameterization num_dept = len(jnp.unique(dept)) z = numpyro.sample('z', dist.Normal(jnp.zeros((num_dept, 2)), 1)) v = jnp.dot(scale_tril, z.T).T logits = v_mu[0] + v[dept, 0] + (v_mu[1] + v[dept, 1]) * male if admit is None: # we use a Delta site to record probs for predictive distribution probs = expit(logits) numpyro.sample('probs', dist.Delta(probs), obs=probs) numpyro.sample('admit', dist.Binomial(applications, logits=logits), obs=admit)
Example #14
Source File: test_examples.py From jaxnet with Apache License 2.0 | 6 votes |
def test_mnist_vae(): @parametrized def encode(input): input = Sequential(Dense(5), relu, Dense(5), relu)(input) mean = Dense(10)(input) variance = Sequential(Dense(10), softplus)(input) return mean, variance decode = Sequential(Dense(5), relu, Dense(5), relu, Dense(5 * 5)) @parametrized def elbo(key, images): mu_z, sigmasq_z = encode(images) logits_x = decode(gaussian_sample(key, mu_z, sigmasq_z)) return bernoulli_logpdf(logits_x, images) - gaussian_kl(mu_z, sigmasq_z) params = elbo.init_parameters(PRNGKey(0), jnp.zeros((32, 5 * 5)), key=PRNGKey(0)) assert (5, 10) == params.encode.sequential1.dense.kernel.shape
Example #15
Source File: util.py From numpyro with Apache License 2.0 | 6 votes |
def _multinomial(key, p, n, n_max, shape=()): if jnp.shape(n) != jnp.shape(p)[:-1]: broadcast_shape = lax.broadcast_shapes(jnp.shape(n), jnp.shape(p)[:-1]) n = jnp.broadcast_to(n, broadcast_shape) p = jnp.broadcast_to(p, broadcast_shape + jnp.shape(p)[-1:]) shape = shape or p.shape[:-1] # get indices from categorical distribution then gather the result indices = categorical(key, p, (n_max,) + shape) # mask out values when counts is heterogeneous if jnp.ndim(n) > 0: mask = promote_shapes(jnp.arange(n_max) < jnp.expand_dims(n, -1), shape=shape + (n_max,))[0] mask = jnp.moveaxis(mask, -1, 0).astype(indices.dtype) excess = jnp.concatenate([jnp.expand_dims(n_max - n, -1), jnp.zeros(jnp.shape(n) + (p.shape[-1] - 1,))], -1) else: mask = 1 excess = 0 # NB: we transpose to move batch shape to the front indices_2D = (jnp.reshape(indices * mask, (n_max, -1,))).T samples_2D = vmap(_scatter_add_one, (0, 0, 0))(jnp.zeros((indices_2D.shape[0], p.shape[-1]), dtype=indices.dtype), jnp.expand_dims(indices_2D, axis=-1), jnp.ones(indices_2D.shape, dtype=indices.dtype)) return jnp.reshape(samples_2D, shape + p.shape[-1:]) - excess
Example #16
Source File: modules.py From jaxnet with Apache License 2.0 | 6 votes |
def GRUCell(carry_size, param_init): @parametrized def gru_cell(carry, x): def param(name): return parameter((x.shape[1] + carry_size, carry_size), param_init, name) both = np.concatenate((x, carry), axis=1) update = sigmoid(np.dot(both, param('update_kernel'))) reset = sigmoid(np.dot(both, param('reset_kernel'))) both_reset_carry = np.concatenate((x, reset * carry), axis=1) compute = np.tanh(np.dot(both_reset_carry, param('compute_kernel'))) out = update * compute + (1 - update) * carry return out, out def carry_init(batch_size): return np.zeros((batch_size, carry_size)) return gru_cell, carry_init
Example #17
Source File: modules.py From jaxnet with Apache License 2.0 | 6 votes |
def BatchNorm(axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True, beta_init=zeros, gamma_init=ones): """Layer construction function for a batch normalization layer.""" axis = (axis,) if np.isscalar(axis) else axis @parametrized def batch_norm(x): ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x))) mean, var = np.mean(x, axis, keepdims=True), fastvar(x, axis, keepdims=True) z = (x - mean) / np.sqrt(var + epsilon) shape = tuple(d for i, d in enumerate(x.shape) if i not in axis) scaled = z * parameter(shape, gamma_init, 'gamma')[ed] if scale else z return scaled + parameter(shape, beta_init, 'beta')[ed] if center else scaled return batch_norm
Example #18
Source File: test_modules.py From jaxnet with Apache License 2.0 | 6 votes |
def test_Regularized(): @parametrized def loss(inputs): a = parameter((), ones, 'a') b = parameter((), lambda key, shape: 2 * jnp.ones(shape), 'b') return a + b reg_loss = Regularized(loss, regularizer=lambda x: x * x) inputs = jnp.zeros(()) params = reg_loss.init_parameters(inputs, key=PRNGKey(0)) assert jnp.array_equal(jnp.ones(()), params.model.a) assert jnp.array_equal(2 * jnp.ones(()), params.model.b) reg_loss_out = reg_loss.apply(params, inputs) assert 1 + 2 + 1 + 4 == reg_loss_out
Example #19
Source File: wavenet.py From jaxnet with Apache License 2.0 | 6 votes |
def Wavenet(dilations, filter_width, initial_filter_width, out_width, residual_channels, dilation_channels, skip_channels, nr_mix): """ :param dilations: dilations for each layer :param filter_width: for the resblock convs :param residual_channels: 1x1 conv output channels :param dilation_channels: gate and filter output channels :param skip_channels: channels before the final output :param initial_filter_width: for the pre processing conv """ @parametrized def wavenet(inputs): hidden = Conv1D(residual_channels, (initial_filter_width,))(inputs) out = np.zeros((hidden.shape[0], out_width, residual_channels), 'float32') for dilation in dilations: res = ResLayer(dilation_channels, residual_channels, filter_width, dilation, out_width)(hidden) hidden, out_partial = res out += out_partial return Sequential(relu, Conv1D(skip_channels, (1,)), relu, Conv1D(3 * nr_mix, (1,)))(out) return wavenet
Example #20
Source File: pixelcnn.py From jaxnet with Apache License 2.0 | 6 votes |
def ConvOrConvTranspose(out_chan, filter_shape=(3, 3), strides=None, padding='SAME', init_scale=1., transpose=False): strides = strides or (1,) * len(filter_shape) def apply(inputs, V, g, b): V = g * _l2_normalize(V, (0, 1, 2)) return (lax.conv_transpose if transpose else _conv)(inputs, V, strides, padding) - b @parametrized 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=jnp.ones(out_chan), b=jnp.zeros(out_chan)) # TODO remove need for `.aval.val` when capturing variables in initializer function: g = Parameter(lambda key: init_scale / jnp.sqrt(jnp.var(example_out.aval.val, (0, 1, 2)) + 1e-10), 'g')() b = Parameter(lambda key: jnp.mean(example_out.aval.val, (0, 1, 2)) * g.aval.val, 'b')() return apply(inputs, V, b, g) return conv_or_conv_transpose
Example #21
Source File: test_modules.py From jaxnet with Apache License 2.0 | 6 votes |
def test_L2Regularized(): @parametrized def loss(inputs): a = parameter((), ones, 'a') b = parameter((), lambda key, shape: 2 * jnp.ones(shape), 'b') return a + b reg_loss = L2Regularized(loss, scale=2) inputs = jnp.zeros(()) params = reg_loss.init_parameters(inputs, key=PRNGKey(0)) assert jnp.array_equal(jnp.ones(()), params.model.a) assert jnp.array_equal(2 * jnp.ones(()), params.model.b) reg_loss_out = reg_loss.apply(params, inputs) assert 1 + 2 + 1 + 4 == reg_loss_out
Example #22
Source File: test_examples.py From jaxnet with Apache License 2.0 | 6 votes |
def test_Parameter_dense(): def Dense(out_dim, kernel_init=glorot_normal(), bias_init=normal()): @parametrized def dense(inputs): kernel = parameter((inputs.shape[-1], out_dim), kernel_init) bias = parameter((out_dim,), bias_init) return jnp.dot(inputs, kernel) + bias return dense net = Dense(2) inputs = jnp.zeros((1, 3)) params = net.init_parameters(inputs, key=PRNGKey(0)) assert (3, 2) == params.parameter0.shape assert (2,) == params.parameter1.shape out = net.apply(params, inputs, jit=True) assert (1, 2) == out.shape
Example #23
Source File: test_examples.py From jaxnet with Apache License 2.0 | 6 votes |
def test_mnist_classifier(): from examples.mnist_classifier import predict, loss, accuracy next_batch = lambda: (jnp.zeros((3, 784)), jnp.zeros((3, 10))) opt = optimizers.Momentum(0.001, mass=0.9) state = opt.init(loss.init_parameters(*next_batch(), key=PRNGKey(0))) t = time.time() for _ in range(10): state = opt.update(loss.apply, state, *next_batch(), jit=True) elapsed = time.time() - t assert 5 > elapsed params = opt.get_parameters(state) train_acc = accuracy.apply_from({loss: params}, *next_batch(), jit=True) assert () == train_acc.shape predict_params = predict.parameters_from({loss.shaped(*next_batch()): params}, next_batch()[0]) predictions = predict.apply(predict_params, next_batch()[0], jit=True) assert (3, 10) == predictions.shape
Example #24
Source File: test_modules.py From jaxnet with Apache License 2.0 | 6 votes |
def test_Batched(): out_dim = 1 @parametrized def unbatched_dense(input): kernel = parameter((out_dim, input.shape[-1]), ones) bias = parameter((out_dim,), ones) return jnp.dot(kernel, input) + bias batch_size = 4 unbatched_params = unbatched_dense.init_parameters(jnp.zeros(2), key=PRNGKey(0)) out = unbatched_dense.apply(unbatched_params, jnp.ones(2)) assert jnp.array([3.]) == out dense_apply = vmap(unbatched_dense.apply, (None, 0)) out_batched_ = dense_apply(unbatched_params, jnp.ones((batch_size, 2))) assert jnp.array_equal(jnp.stack([out] * batch_size), out_batched_) dense = Batched(unbatched_dense) params = dense.init_parameters(jnp.ones((batch_size, 2)), key=PRNGKey(0)) assert_parameters_equal((unbatched_params,), params) out_batched = dense.apply(params, jnp.ones((batch_size, 2))) assert jnp.array_equal(out_batched_, out_batched)
Example #25
Source File: jax.py From deepx with MIT License | 5 votes |
def block_sum(self, X, m, n): leading_dim = self.shape(X)[:-2] block_sum = self.zeros(self.concat([leading_dim, [m, m]], 0)) for i in range(n): block_sum += X[..., i*m:(i+1)*m, i*m:(i+1)*m] return block_sum
Example #26
Source File: jax.py From deepx with MIT License | 5 votes |
def zeros(self, shape, dtype=None, name=None): dtype = dtype or self.floatx() if not isinstance(shape, int): shape = tuple(shape) return np.zeros(shape, dtype=dtype)
Example #27
Source File: test_modules.py From jaxnet with Apache License 2.0 | 5 votes |
def test_GRUCell_shape(): gru_cell, init_carry = GRUCell(10, zeros) x = jnp.zeros((2, 3)) carry = init_carry(batch_size=2) params = gru_cell.init_parameters(carry, x, key=PRNGKey(0)) out = gru_cell.apply(params, carry, x) assert (2, 10) == out[0].shape assert (2, 10) == out[1].shape
Example #28
Source File: test_modules.py From jaxnet with Apache License 2.0 | 5 votes |
def test_Reparametrized(): @parametrized def net(inputs): return parameter((), lambda key, shape: 2 * jnp.ones(shape)) scaled_net = Reparametrized(net, reparametrization_factory=Scaled) inputs = jnp.zeros(()) params = scaled_net.init_parameters(inputs, key=PRNGKey(0)) reg_loss_out = scaled_net.apply(params, inputs) assert 4 == reg_loss_out
Example #29
Source File: agent.py From bsuite with Apache License 2.0 | 5 votes |
def default_agent(obs_spec: specs.Array, action_spec: specs.DiscreteArray, seed: int = 0) -> base.Agent: """Creates an actor-critic agent with default hyperparameters.""" hidden_size = 256 initial_rnn_state = hk.LSTMState( hidden=jnp.zeros((1, hidden_size), dtype=jnp.float32), cell=jnp.zeros((1, hidden_size), dtype=jnp.float32)) def network(inputs: jnp.ndarray, state) -> Tuple[Tuple[Logits, Value], LSTMState]: flat_inputs = hk.Flatten()(inputs) torso = hk.nets.MLP([hidden_size, hidden_size]) lstm = hk.LSTM(hidden_size) policy_head = hk.Linear(action_spec.num_values) value_head = hk.Linear(1) embedding = torso(flat_inputs) embedding, state = lstm(embedding, state) logits = policy_head(embedding) value = value_head(embedding) return (logits, jnp.squeeze(value, axis=-1)), state return ActorCriticRNN( obs_spec=obs_spec, action_spec=action_spec, network=network, initial_rnn_state=initial_rnn_state, optimizer=optix.adam(3e-3), rng=hk.PRNGSequence(seed), sequence_length=32, discount=0.99, td_lambda=0.9, )
Example #30
Source File: funnel.py From numpyro with Apache License 2.0 | 5 votes |
def model(dim=10): y = numpyro.sample('y', dist.Normal(0, 3)) numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))