AttributeError: partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)

[limin321]

I installed cell2location in a new conda environment following the instruction.
Here is my jax and jaxlib version:

(cell2loc_env) limin@DERM-6NBQ05N-LT jupyterNotebook % pip show jax
Name: jax
Version: 0.3.23
Summary: Differentiate, compile, and transform Numpy code.
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: [email protected]
License: Apache-2.0
Location: /Users/limin/miniconda3/envs/cell2loc_env/lib/python3.9/site-packages
Requires: absl-py, etils, numpy, opt-einsum, scipy, typing-extensions
Required-by: chex, flax, numpyro, optax, orbax, scvi-tools
(cell2loc_env) limin@DERM-6NBQ05N-LT jupyterNotebook % pip show jaxlib
Name: jaxlib
Version: 0.3.22
Summary: XLA library for JAX
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: [email protected]
License: Apache-2.0
Location: /Users/limin/miniconda3/envs/cell2loc_env/lib/python3.9/site-packages
Requires: absl-py, numpy, scipy
Required-by: chex, numpyro, optax, orbax, scvi-tools
(cell2loc_env) limin@DERM-6NBQ05N-LT jupyterNotebook % 

When I ran the following codes:

import cell2location
import scvi

I ran into this error:

I tried several versions of jax and jaxlib. Always giving me the same error.
Any suggestion would be very helpful.
My mac is Apple M1 chip.

Thank you so much.

Best,
LC

[vitkl]

Hi @limin321

You can try this #133 (comment)

M1/M2 Macs have issues with various packages.

Maybe this also helps:

conda create -y -n schierarchy-env python=3.8
conda activate schierarchy-env
pip install git+https://github.com/pyro-ppl/pyro.git@dev
conda install -y -c anaconda hdf5 pytables netcdf4
pip install git+https://github.com/BayraktarLab/cell2location.git#egg=cell2location[tutorials]

If you find a solution that works, please post it here.

[vitkl]

There is also this discussion #206 (comment)

[limin321]

Hi @vitkl
I tried your second option by running the following script, with no errors or warnings. I assume this works.

conda create -y -n schierarchy-env python=3.8
conda activate schierarchy-env
pip install git+https://github.com/pyro-ppl/pyro.git@dev
conda install -y -c anaconda hdf5 pytables netcdf4
pip install git+https://github.com/BayraktarLab/cell2location.git#egg=cell2location[tutorials]

However, when I tried to train the model, I got an error again. Because I run the code in the local Mac, which means I don't have GPU,so I set the environment variable as follows and with error:

import os
myintvariable = 1
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = str(myintvariable)
mod.train(max_epochs=250, use_gpu=True)

GPU available: True (mps), used: True TPU available: False, using: 0 TPU cores IPU available: False, using: 0 IPUs HPU available: False, using: 0 HPUs /Users/limin/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pytorch_lightning/trainer/configuration_validator.py:105: UserWarning: You passed in a val_dataloaderbut have novalidation_step. Skipping val loop. rank_zero_warn("You passed in a val_dataloaderbut have novalidation_step`. Skipping val loop.")

AttributeError Traceback (most recent call last)
Input In [76], in <cell line: 4>()
2 myintvariable = 1
3 os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = str(myintvariable)
----> 4 mod.train(max_epochs=250, use_gpu=True)

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/cell2location/models/reference/_reference_model.py:157, in RegressionModel.train(self, max_epochs, batch_size, train_size, lr, **kwargs)
154 kwargs["train_size"] = train_size
155 kwargs["lr"] = lr
--> 157 super().train(**kwargs)

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/scvi/model/base/_pyromixin.py:146, in PyroSviTrainMixin.train(self, max_epochs, use_gpu, train_size, validation_size, batch_size, early_stopping, lr, training_plan, plan_kwargs, **trainer_kwargs)
136 trainer_kwargs["callbacks"].append(PyroJitGuideWarmup())
138 runner = TrainRunner(
139 self,
140 training_plan=training_plan,
(...)
144 **trainer_kwargs,
145 )
--> 146 return runner()

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/scvi/train/_trainrunner.py:82, in TrainRunner.call(self)
79 if hasattr(self.data_splitter, "n_val"):
80 self.training_plan.n_obs_validation = self.data_splitter.n_val
---> 82 self.trainer.fit(self.training_plan, self.data_splitter)
83 self._update_history()
85 # data splitter only gets these attrs after fit

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/scvi/train/_trainer.py:188, in Trainer.fit(self, *args, **kwargs)
182 if isinstance(args[0], PyroTrainingPlan):
183 warnings.filterwarnings(
184 action="ignore",
185 category=UserWarning,
186 message="LightningModule.configure_optimizers returned None",
187 )
--> 188 super().fit(*args, **kwargs)

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:696, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
677 r"""
678 Runs the full optimization routine.
679
(...)
693 datamodule: An instance of :class:~pytorch_lightning.core.datamodule.LightningDataModule.
694 """
695 self.strategy.model = model
--> 696 self._call_and_handle_interrupt(
697 self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
698 )

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:650, in Trainer._call_and_handle_interrupt(self, trainer_fn, *args, **kwargs)
648 return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
649 else:
--> 650 return trainer_fn(*args, **kwargs)
651 # TODO(awaelchli): Unify both exceptions below, where KeyboardError doesn't re-raise
652 except KeyboardInterrupt as exception:

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:735, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
731 ckpt_path = ckpt_path or self.resume_from_checkpoint
732 self._ckpt_path = self.__set_ckpt_path(
733 ckpt_path, model_provided=True, model_connected=self.lightning_module is not None
734 )
--> 735 results = self._run(model, ckpt_path=self.ckpt_path)
737 assert self.state.stopped
738 self.training = False

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1166, in Trainer._run(self, model, ckpt_path)
1162 self._checkpoint_connector.restore_training_state()
1164 self._checkpoint_connector.resume_end()
-> 1166 results = self._run_stage()
1168 log.detail(f"{self.class.name}: trainer tearing down")
1169 self._teardown()

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1252, in Trainer._run_stage(self)
1250 if self.predicting:
1251 return self._run_predict()
-> 1252 return self._run_train()

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1283, in Trainer._run_train(self)
1280 self.fit_loop.trainer = self
1282 with torch.autograd.set_detect_anomaly(self._detect_anomaly):
-> 1283 self.fit_loop.run()

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py:195, in Loop.run(self, *args, **kwargs)
191 return self.on_skip()
193 self.reset()
--> 195 self.on_run_start(*args, **kwargs)
197 while not self.done:
198 try:

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py:222, in FitLoop.on_run_start(self)
219 self._is_fresh_start_epoch = True
220 self._results.to(device=self.trainer.lightning_module.device)
--> 222 self.trainer._call_callback_hooks("on_train_start")
223 self.trainer._call_lightning_module_hook("on_train_start")
224 self.trainer._call_strategy_hook("on_train_start")

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1597, in Trainer._call_callback_hooks(self, hook_name, *args, **kwargs)
1595 if callable(fn):
1596 with self.profiler.profile(f"[Callback]{callback.state_key}.{hook_name}"):
-> 1597 fn(self, self.lightning_module, *args, **kwargs)
1599 if pl_module:
1600 # restore current_fx when nested context
1601 pl_module._current_fx_name = prev_fx_name

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/scvi/model/base/_pyromixin.py:45, in PyroJitGuideWarmup.on_train_start(self, trainer, pl_module)
43 tens = {k: t.to(pl_module.device) for k, t in tensors.items()}
44 args, kwargs = pl_module.module._get_fn_args_from_batch(tens)
---> 45 pyro_guide(*args, **kwargs)
46 break

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pyro/nn/module.py:448, in PyroModule.call(self, *args, **kwargs)
446 def call(self, *args, **kwargs):
447 with self._pyro_context:
--> 448 result = super().call(*args, **kwargs)
449 if (
450 pyro.settings.get("validate_poutine")
451 and not self._pyro_context.active
452 and _is_module_local_param_enabled()
453 ):
454 self._check_module_local_param_usage()

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
1190 # If we don't have any hooks, we want to skip the rest of the logic in
1191 # this function, and just call forward.
1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1193 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194 return forward_call(*input, **kwargs)
1195 # Do not call functions when jit is used
1196 full_backward_hooks, non_full_backward_hooks = [], []

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pyro/infer/autoguide/guides.py:522, in AutoNormal.forward(self, *args, **kwargs)
519 if frame.vectorized:
520 stack.enter_context(plates[frame.name])
--> 522 site_loc, site_scale = self._get_loc_and_scale(name)
523 unconstrained_latent = pyro.sample(
524 name + "_unconstrained",
525 dist.Normal(
(...)
529 infer={"is_auxiliary": True},
530 )
532 value = transform(unconstrained_latent)

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pyro/infer/autoguide/guides.py:495, in AutoNormal._get_loc_and_scale(self, name)
493 def _get_loc_and_scale(self, name):
494 site_loc = deep_getattr(self.locs, name)
--> 495 site_scale = deep_getattr(self.scales, name)
496 return site_loc, site_scale

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pyro/infer/autoguide/utils.py:23, in deep_getattr(obj, key)
21 def deep_getattr(obj, key):
22 for part in key.split("."):
---> 23 obj = getattr(obj, part)
24 return obj

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pyro/nn/module.py:475, in PyroModule.getattr(self, name)
473 if name in _pyro_params:
474 constraint, event_dim = _pyro_params[name]
--> 475 unconstrained_value = getattr(self, name + "_unconstrained")
476 if self._pyro_context.active:
477 fullname = self._pyro_get_fullname(name)

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/pyro/nn/module.py:528, in PyroModule.getattr(self, name)
525 prior = prior(self)
526 return prior()
--> 528 result = super().getattr(name)
530 # Regular nn.Parameters trigger pyro.param statements.
531 if isinstance(result, torch.nn.Parameter) and not name.endswith(
532 "_unconstrained"
533 ):

File ~/miniconda3/envs/schierarchy-env/lib/python3.8/site-packages/torch/nn/modules/module.py:1269, in Module.getattr(self, name)
1267 if name in modules:
1268 return modules[name]
-> 1269 raise AttributeError("'{}' object has no attribute '{}'".format(
1270 type(self).name, name))

AttributeError: 'PyroModule' object has no attribute 'per_cluster_mu_fg_unconstrained'`

Any suggestion on how to solve this issue?

Best,
LC

[limin321]

I made it run on HPC with GPU. I would like to close my question. Thank you for all the help.