def debug_vmap_bfgs():
import jax.numpy as jnp
from jax import jit, config
from jax.scipy.optimize import minimize
import os
config.enable_omnistaging()
ncpu=2
os.environ['XLA_FLAGS'] = f"--xla_force_host_platform_device_count={ncpu}"
def cost_fn(x):
return -jnp.sum(x**2)
x = random.uniform(random.PRNGKey(0), (3,), minval=-1, maxval=1)
result = jit(lambda x: minimize(cost_fn, x, method='BFGS'))(x)
print(result)
jax.tree_util.tree_leaves(new_dist)):
self.assertEqual(type(old_param), type(new_param))
if isinstance(old_param, np.ndarray):
self.assertTupleEqual(old_param.shape, new_param.shape)
@test_all_distributions
@hp.given(hps.data())
@tfp_hps.tfp_hp_settings(default_max_examples=DEFAULT_MAX_EXAMPLES)
def testInputOutputOfJittedFunction(self, dist_name, data):
if (dist_name in PYTREE_BLOCKLIST) != FLAGS.blocklists_only:
self.skipTest('Distribution currently broken.')
@jax.jit
def dist_and_sample(dist):
return dist, dist.sample(seed=test_util.test_seed())
dist = data.draw(
dhps.distributions(enable_vars=False,
dist_name=dist_name,
validate_args=False,
eligibility_filter=lambda dname: dname not in
PYTREE_BLOCKLIST))
dist_and_sample(dist)
if __name__ == '__main__':
config.enable_omnistaging()
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'
tf.test.main()
# flax
import flax
from flax.core import freeze, unfreeze
from flax import linen as nn
# constants
LARGE_NEG_VALUE = -1e10
# config
from jax.config import config
config.enable_omnistaging() # Linen requires enabling omnistaging
# helpers
def cross_entropy(logits, targets, axis=-1):
logprobs = nn.log_softmax(logits, axis=axis)
nll = np.take_along_axis(logprobs,
np.expand_dims(targets, axis=axis),
axis=axis)
ce = -np.mean(nll)
return ce
def fixed_pos_embedding(seq, dim):
inv_freq = 1.0 / (10000**(np.arange(0, dim, 2) / dim))
from jax.config import config as jax_config
from jax import random
import jax.numpy as np
from jax import ops
from jax.api import grad
from jax import test_util as jtu
from jax import jit, vmap
from jax_md import smap, partition, space, energy, quantity
from jax_md.util import *
jax_config.parse_flags_with_absl()
jax_config.enable_omnistaging()
FLAGS = jax_config.FLAGS
PARTICLE_COUNT = 1000
STOCHASTIC_SAMPLES = 10
SPATIAL_DIMENSION = [2, 3]
if FLAGS.jax_enable_x64:
POSITION_DTYPE = [f32, f64]
else:
POSITION_DTYPE = [f32]
class CellListTest(jtu.JaxTestCase):
@parameterized.named_parameters(
jtu.cases_from_list({
def setUp(self): super().setUp() config.enable_omnistaging()