avoid math.prod for python 3.7

Summary: This makes the new volumes tutorial work on google colab.

Reviewed By: kjchalup

Differential Revision: D38501906

fbshipit-source-id: a606a357e929dae903dc4d9067bd1519f05b1458

pytorch3d/common/compat.py

@@ -10,7 +10,7 @@ import torch
 
 
 """
-Some functions which depend on PyTorch versions.
+Some functions which depend on PyTorch or Python versions.
 """
 
 
@@ -79,3 +79,12 @@ def meshgrid_ij(
     # pyre-fixme[6]: For 1st param expected `Union[List[Tensor], Tensor]` but got
     #  `Union[Sequence[Tensor], Tensor]`.
     return torch.meshgrid(*A)
+
+
+def prod(iterable, *, start=1):
+    """
+    Like math.prod in Python 3.8 and later.
+    """
+    for i in iterable:
+        start *= i
+    return start

pytorch3d/implicitron/models/generic_model.py

@@ -17,6 +17,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import torch
 import tqdm
 from omegaconf import DictConfig
+from pytorch3d.common.compat import prod
 from pytorch3d.implicitron.models.metrics import (
     RegularizationMetricsBase,
     ViewMetricsBase,
@@ -919,7 +920,7 @@ def _chunk_generator(
             f"by n_pts_per_ray ({n_pts_per_ray})"
         )
 
-    n_rays = math.prod(spatial_dim)
+    n_rays = prod(spatial_dim)
     # special handling for raytracing-based methods
     n_chunks = -(-n_rays * max(n_pts_per_ray, 1) // chunk_size)
     chunk_size_in_rays = -(-n_rays // n_chunks)
@@ -935,9 +936,9 @@ def _chunk_generator(
             directions=ray_bundle.directions.reshape(batch_size, -1, 3)[
                 :, start_idx:end_idx
             ],
-            lengths=ray_bundle.lengths.reshape(
-                batch_size, math.prod(spatial_dim), n_pts_per_ray
-            )[:, start_idx:end_idx],
+            lengths=ray_bundle.lengths.reshape(batch_size, n_rays, n_pts_per_ray)[
+                :, start_idx:end_idx
+            ],
             xys=ray_bundle.xys.reshape(batch_size, -1, 2)[:, start_idx:end_idx],
         )
         extra_args = kwargs.copy()

pytorch3d/implicitron/models/implicit_function/utils.py

@@ -4,10 +4,10 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import math
 from typing import Callable, Optional
 
 import torch
+from pytorch3d.common.compat import prod
 from pytorch3d.renderer.cameras import CamerasBase
 
 
@@ -52,7 +52,7 @@ def create_embeddings_for_implicit_function(
         embeds = torch.empty(
             bs,
             1,
-            math.prod(spatial_size),
+            prod(spatial_size),
             pts_per_ray,
             0,
             dtype=xyz_world.dtype,
@@ -62,7 +62,7 @@ def create_embeddings_for_implicit_function(
         embeds = xyz_embedding_function(ray_points_for_embed).reshape(
             bs,
             1,
-            math.prod(spatial_size),
+            prod(spatial_size),
             pts_per_ray,
             -1,
         )  # flatten spatial, add n_src dim
@@ -73,7 +73,7 @@ def create_embeddings_for_implicit_function(
         embed_shape = (
             bs,
             embeds_viewpooled.shape[1],
-            math.prod(spatial_size),
+            prod(spatial_size),
             pts_per_ray,
             -1,
         )

pytorch3d/implicitron/models/renderer/sdf_renderer.py

@@ -3,11 +3,11 @@
 #              implicit_differentiable_renderer.py
 # Copyright (c) 2020 Lior Yariv
 import functools
-import math
 from typing import List, Optional, Tuple
 
 import torch
 from omegaconf import DictConfig
+from pytorch3d.common.compat import prod
 from pytorch3d.implicitron.tools.config import (
     get_default_args_field,
     registry,
@@ -105,7 +105,7 @@ class SignedDistanceFunctionRenderer(BaseRenderer, torch.nn.Module):  # pyre-ign
 
         # object_mask: silhouette of the object
         batch_size, *spatial_size, _ = ray_bundle.lengths.shape
-        num_pixels = math.prod(spatial_size)
+        num_pixels = prod(spatial_size)
 
         cam_loc = ray_bundle.origins.reshape(batch_size, -1, 3)
         ray_dirs = ray_bundle.directions.reshape(batch_size, -1, 3)