Python jax.numpy.zeros_like() Examples
The following are 4
code examples of jax.numpy.zeros_like().
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: test_reparam.py From numpyro with Apache License 2.0 | 6 votes |
def test_log_normal(shape): loc = np.random.rand(*shape) * 2 - 1 scale = np.random.rand(*shape) + 0.5 def model(): with numpyro.plate_stack("plates", shape): with numpyro.plate("particles", 100000): return numpyro.sample("x", dist.TransformedDistribution( dist.Normal(jnp.zeros_like(loc), jnp.ones_like(scale)), [AffineTransform(loc, scale), ExpTransform()]).expand_by([100000])) with handlers.trace() as tr: value = handlers.seed(model, 0)() expected_moments = get_moments(value) with numpyro.handlers.reparam(config={"x": TransformReparam()}): with handlers.trace() as tr: value = handlers.seed(model, 0)() assert tr["x"]["type"] == "deterministic" actual_moments = get_moments(value) assert_allclose(actual_moments, expected_moments, atol=0.05)
Example #2
Source File: hessian_computation.py From spectral-density with Apache License 2.0 | 5 votes |
def _tree_zeros_like(tree): def f(x): return np.zeros_like(x) return tu.tree_map(f, tree)
Example #3
Source File: jax.py From deepx with MIT License | 5 votes |
def zeros_like(self, x, dtype=None, name=None): return np.zeros_like(x, dtype=dtype)
Example #4
Source File: tke_jax.py From pyhpc-benchmarks with The Unlicense | 5 votes |
def adv_flux_superbee_wgrid(var, u_wgrid, v_wgrid, w_wgrid, maskW, dxt, dyt, dzw, cost, cosu, dt_tracer): """ Calculates advection of a tracer defined on Wgrid """ maskUtr = np.zeros_like(maskW) maskUtr = jax.ops.index_update( maskUtr, jax.ops.index[:-1, :, :], maskW[1:, :, :] * maskW[:-1, :, :] ) adv_fe = np.zeros_like(maskW) adv_fe = jax.ops.index_update( adv_fe, jax.ops.index[1:-2, 2:-2, :], _adv_superbee(u_wgrid, var, maskUtr, dxt, 0, cost, cosu, dt_tracer) ) maskVtr = np.zeros_like(maskW) maskVtr = jax.ops.index_update( maskVtr, jax.ops.index[:, :-1, :], maskW[:, 1:, :] * maskW[:, :-1, :] ) adv_fn = np.zeros_like(maskW) adv_fn = jax.ops.index_update( adv_fn, jax.ops.index[2:-2, 1:-2, :], _adv_superbee(v_wgrid, var, maskVtr, dyt, 1, cost, cosu, dt_tracer) ) maskWtr = np.zeros_like(maskW) maskWtr = jax.ops.index_update( maskWtr, jax.ops.index[:, :, :-1], maskW[:, :, 1:] * maskW[:, :, :-1] ) adv_ft = np.zeros_like(maskW) adv_ft = jax.ops.index_update( adv_ft, jax.ops.index[2:-2, 2:-2, :-1], _adv_superbee(w_wgrid, var, maskWtr, dzw, 2, cost, cosu, dt_tracer) ) return adv_fe, adv_fn, adv_ft