google-deepmind / enn Goto Github PK

AttributeError                            Traceback (most recent call last)
Cell In[5], line 25
     17 enn = networks.MLPEnsembleMatchedPrior(
     18     output_sizes=[50, 50, 1],
     19     num_ensemble=10,
     20     dummy_input=np.zeros(50)
     21 )
     23 # Loss
     24 loss_fn = losses.average_single_index_loss(
---> 25     single_loss=losses.L2LossWithBootstrap(),
     26     num_index_samples=10
     27 )
     29 # Optimizer
     30 optimizer = optax.adam(1e-3)

AttributeError: module 'enn.losses' has no attribute 'L2LossWithBootstrap'

from enn.loggers import TerminalLogger

from enn import losses
from enn import networks
from enn import supervised
from enn.supervised import regression_data
import optax
import numpy as np

# A small dummy dataset
dataset = regression_data.make_dataset()

# Logger
logger = TerminalLogger('supervised_regression')

# ENN
enn = networks.MLPEnsembleMatchedPrior(
    output_sizes=[50, 50, 1],
    num_ensemble=10,
    dummy_input=np.zeros(50)
)

# Loss
loss_fn = losses.average_single_index_loss(
    single_loss=losses.L2LossWithBootstrap(),
    num_index_samples=10
)

# Optimizer
optimizer = optax.adam(1e-3)

# Train the experiment
experiment = supervised.Experiment(
    enn, loss_fn, optimizer, dataset, seed=0, logger=logger)
experiment.train(FLAGS.num_batch)

    dummy_input=np.zeros(50)

# Train the experiment
experiment.train(FLAGS.num_batch)

AttributeError                            Traceback (most recent call last)
[/usr/local/lib/python3.8/dist-packages/enn/networks/ensembles.py](https://kh9bbgsdon-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20221220-060108-RC02_496713401#) in apply(params, states, inputs, index)
     82       sub_states = jax.tree_map(particle_selector, states)
     83       out, new_sub_states = model.apply(sub_params, sub_states, inputs)
---> 84       new_states = jax.tree_multimap(
     85           lambda s, nss: s.at[index, ...].set(nss), states, new_sub_states)
     86       return out, new_states

AttributeError: module 'jax' has no attribute 'tree_multimap'

AttributeError                            Traceback (most recent call last)
<ipython-input-7-4863fdbb2553> in <module>()
     13     num_ensemble=FLAGS.index_dim,
     14     prior_scale=FLAGS.prior_scale,
---> 15     seed=FLAGS.seed,
     16 )
     17 

4 frames
/usr/local/lib/python3.7/dist-packages/enn/networks/ensembles.py in __init__(self, output_sizes, dummy_input, num_ensemble, prior_scale, seed, w_init, b_init)
    137     """Ensemble of MLPs with matched prior functions."""
    138     mlp_priors = make_mlp_ensemble_prior_fns(
--> 139         output_sizes, dummy_input, num_ensemble, seed)
    140     enn = priors.EnnWithAdditivePrior(
    141         enn=MLPEnsembleEnn(

/usr/local/lib/python3.7/dist-packages/enn/networks/ensembles.py in make_mlp_ensemble_prior_fns(output_sizes, dummy_input, num_ensemble, seed, w_init, b_init)
     90     return hk.Sequential(layers)(x)
     91 
---> 92   transformed = hk.without_apply_rng(hk.transform(net_fn))
     93 
     94   prior_fns = []

/usr/local/lib/python3.7/dist-packages/haiku/_src/transform.py in transform(f, apply_rng)
    301         "Replace hk.transform(..., apply_rng=True) with hk.transform(...).")
    302 
--> 303   return without_state(transform_with_state(f))
    304 
    305 

/usr/local/lib/python3.7/dist-packages/haiku/_src/transform.py in transform_with_state(f)
    359   """
    360   analytics.log_once("transform_with_state")
--> 361   check_not_jax_transformed(f)
    362 
    363   unexpected_tracer_hint = (

/usr/local/lib/python3.7/dist-packages/haiku/_src/transform.py in check_not_jax_transformed(f)
    306 def check_not_jax_transformed(f):
    307   # TODO(tomhennigan): Consider `CompiledFunction = type(jax.jit(lambda: 0))`.
--> 308   if isinstance(f, (jax.xla.xe.CompiledFunction, jax.xla.xe.PmapFunction)):  # pytype: disable=name-error
    309     raise ValueError("A common error with Haiku is to pass an already jit "
    310                      "(or pmap) decorated function into hk.transform (e.g. "

AttributeError: module 'jaxlib.xla_extension' has no attribute 'PmapFunction'