V2

nerfies/configs.py

@@ -16,7 +16,7 @@
 from typing import Any, Mapping, Optional, Tuple
 
 import dataclasses
-from flax import nn
+from flax import linen as nn
 import gin
 import immutabledict
 

nerfies/evaluation.py

@@ -92,7 +92,7 @@ def render_image(
     ret_map = jax_utils.unreplicate(model_out[ret_key])
     ret_map = jax.tree_map(lambda x: utils.unshard(x, padding), ret_map)
     ret_maps.append(ret_map)
-  ret_map = jax.tree_multimap(lambda *x: jnp.concatenate(x, axis=0), *ret_maps)
+  ret_map = jax.tree_map(lambda *x: jnp.concatenate(x, axis=0), *ret_maps)
   logging.info('Rendering took %.04s', time.time() - start_time)
   out = {}
   for key, value in ret_map.items():

nerfies/model_utils.py

@@ -104,7 +104,7 @@ def volumetric_rendering(rgb,
   last_sample_z = 1e10 if sample_at_infinity else 1e-19
   dists = jnp.concatenate([
       z_vals[..., 1:] - z_vals[..., :-1],
-      jnp.broadcast_to([last_sample_z], z_vals[..., :1].shape)
+      jnp.broadcast_to(jnp.array([last_sample_z]), z_vals[..., :1].shape)
   ], -1)
   dists = dists * jnp.linalg.norm(dirs[..., None, :], axis=-1)
   alpha = 1.0 - jnp.exp(-sigma * dists)

nerfies/schedules.py

@@ -15,6 +15,7 @@
 """Annealing Schedules."""
 import abc
 import collections
+from collections.abc import Mapping
 import copy
 import math
 from typing import Any, Iterable, List, Tuple, Union
@@ -38,7 +39,7 @@ def from_config(schedule):
     return schedule
   if isinstance(schedule, Tuple) or isinstance(schedule, List):
     return from_tuple(schedule)
-  if isinstance(schedule, collections.Mapping):
+  if isinstance(schedule, Mapping):
     return from_dict(schedule)
 
   raise ValueError(f'Unknown type {type(schedule)}.')

nerfies/utils.py

@@ -377,7 +377,7 @@ def strided_subset(sequence, count):
 
 def tree_collate(list_of_pytrees):
   """Collates a list of pytrees with the same structure."""
-  return tree_util.tree_multimap(lambda *x: np.stack(x), *list_of_pytrees)
+  return tree_util.tree_map(lambda *x: np.stack(x), *list_of_pytrees)
 
 
 @contextlib.contextmanager