fix axis arg to dim for pytorch funcs

pyrad/cameras/rays.py

@@ -117,7 +117,7 @@ class RaySamples(TensorDataclass):
         # mip-nerf version of transmittance calculation:
         transmittance = torch.cumsum(delta_density[..., :-1, :], dim=-2)
         transmittance = torch.cat(
-            [torch.zeros((*transmittance.shape[:1], 1, 1)).to(densities.device), transmittance], axis=-2
+            [torch.zeros((*transmittance.shape[:1], 1, 1)).to(densities.device), transmittance], dim=-2
         )
         transmittance = torch.exp(-transmittance)  # [..., "num_samples"]
 

pyrad/fields/modules/encoding.py

@@ -18,6 +18,7 @@ Encoding functions
 
 from abc import abstractmethod
 from typing import Optional
+
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -132,7 +133,7 @@ class NeRFEncoding(Encoding):
             )
 
         if self.include_input:
-            encoded_inputs = torch.cat([encoded_inputs, in_tensor], axis=-1)
+            encoded_inputs = torch.cat([encoded_inputs, in_tensor], dim=-1)
         return encoded_inputs
 
 
@@ -185,7 +186,7 @@ class RFFEncoding(Encoding):
             )
 
         if self.include_input:
-            encoded_inputs = torch.cat([encoded_inputs, in_tensor], axis=-1)
+            encoded_inputs = torch.cat([encoded_inputs, in_tensor], dim=-1)
 
         return encoded_inputs
 

pyrad/graphs/modules/ray_sampler.py

@@ -313,7 +313,7 @@ class PDFSampler(Sampler):
                 (ray_samples.frustums.starts[..., 1:, 0] + ray_samples.frustums.ends[..., :-1, 0]) / 2.0,
                 ray_samples.frustums.ends[..., -1:, 0],
             ],
-            axis=-1,
+            dim=-1,
         )
 
         inds = torch.searchsorted(cdf, u, side="right")

pyrad/utils/math.py

@@ -15,6 +15,7 @@
 """ Math Helper Functions """
 
 from dataclasses import dataclass
+
 import torch
 from torchtyping import TensorType
 
@@ -117,7 +118,7 @@ def compute_3d_gaussian(
 
     dir_outer_product = directions[..., :, None] * directions[..., None, :]
     eye = torch.eye(directions.shape[-1], device=directions.device)
-    dir_mag_sq = torch.clamp(torch.sum(directions**2, axis=-1, keepdim=True), min=1e-10)
+    dir_mag_sq = torch.clamp(torch.sum(directions**2, dim=-1, keepdim=True), min=1e-10)
     null_outer_product = eye - directions[..., :, None] * (directions / dir_mag_sq)[..., None, :]
     dir_cov_diag = dir_variance[..., None] * dir_outer_product[..., :, :]
     radius_cov_diag = radius_variance[..., None] * null_outer_product[..., :, :]

pyrad/utils/plotly.py

@@ -171,7 +171,7 @@ def get_sphere(
     x = torch.cos(theta) * torch.sin(phi)
     y = torch.cos(theta) * torch.cos(phi)
     z = torch.sin(theta)
-    pts = torch.stack((x, y, z), axis=-1)
+    pts = torch.stack((x, y, z), dim=-1)
 
     pts *= radius
     if center is not None:
@@ -220,10 +220,10 @@ def get_gaussian_ellipsiod(
     x = torch.cos(theta) * torch.sin(phi)
     y = torch.cos(theta) * torch.cos(phi)
     z = torch.sin(theta)
-    pts = torch.stack((x, y, z), axis=-1)
+    pts = torch.stack((x, y, z), dim=-1)
 
     eigenvals, eigenvecs = torch.linalg.eigh(cov)
-    idx = torch.sum(cov, axis=0).argsort()
+    idx = torch.sum(cov, dim=0).argsort()
     idx = eigenvals[idx].argsort()
     eigenvals = eigenvals[idx]
     eigenvecs = eigenvecs[:, idx]

tests/renderers/test_renderers.py

@@ -3,8 +3,8 @@ Test renderers
 """
 import pytest
 import torch
-from pyrad.cameras.rays import Frustums, RaySamples
 
+from pyrad.cameras.rays import Frustums, RaySamples
 from pyrad.renderers import renderers
 
 
@@ -14,7 +14,7 @@ def test_rgb_renderer():
 
     rgb_samples = torch.ones((3, num_samples, 3))
     weights = torch.ones((3, num_samples, 1))
-    weights /= torch.sum(weights, axis=-2, keepdim=True)
+    weights /= torch.sum(weights, dim=-2, keepdim=True)
 
     rgb_renderer = renderers.RGBRenderer()
 
@@ -33,7 +33,7 @@ def test_sh_renderer():
 
     sh = torch.ones((3, num_samples, 3 * levels**2))
     weights = torch.ones((3, num_samples, 1))
-    weights /= torch.sum(weights, axis=-2, keepdim=True)
+    weights /= torch.sum(weights, dim=-2, keepdim=True)
     directions = torch.zeros((3, num_samples, 3))
     directions[..., 0] = 1
 
@@ -48,7 +48,7 @@ def test_acc_renderer():
 
     num_samples = 10
     weights = torch.ones((3, num_samples, 1))
-    weights /= torch.sum(weights, axis=-2, keepdim=True)
+    weights /= torch.sum(weights, dim=-2, keepdim=True)
 
     acc_renderer = renderers.AccumulationRenderer()
 
@@ -61,7 +61,7 @@ def test_depth_renderer():
 
     num_samples = 10
     weights = torch.ones((3, num_samples, 1))
-    weights /= torch.sum(weights, axis=-2, keepdim=True)
+    weights /= torch.sum(weights, dim=-2, keepdim=True)
 
     frustums = Frustums.get_mock_frustum()
     frustums.starts = torch.linspace(0, 100, num_samples)[..., None]