Stricter typing, update docstrings (#226)

nerfactory/cameras/rays.py

@@ -31,19 +31,19 @@ from nerfactory.utils.tensor_dataclass import TensorDataclass
 class Frustums(TensorDataclass):
     """Describes region of space as a frustum.
 
-    Args:
-        origins (TensorType[..., 3]): xyz coordinate for ray origin.
-        directions (TensorType[..., 3]): Direction of ray.
-        starts (TensorType[..., num_samples, 1]): Where the frustum starts along a ray.
-        ends (TensorType[..., num_samples, 1]): Where the frustum ends along a ray.
-        pixel_area (TensorType[..., 1]): Projected area of pixel a distance 1 away from origin.
+    Attributes:
+        origins: xyz coordinate for ray origin.
+        directions: Direction of ray.
+        starts: Where the frustum starts along a ray.
+        ends: Where the frustum ends along a ray.
+        pixel_area: Projected area of pixel a distance 1 away from origin.
     """
 
-    origins: TensorType[..., 3]
-    directions: TensorType[..., 3]
-    starts: TensorType[..., 1]
-    ends: TensorType[..., 1]
-    pixel_area: TensorType[..., 1]
+    origins: TensorType["bs":..., 3]
+    directions: TensorType["bs":..., 3]
+    starts: TensorType["bs":..., 1]
+    ends: TensorType["bs":..., 1]
+    pixel_area: TensorType["bs":..., 1]
 
     def get_positions(self) -> TensorType[..., 3]:
         """Calulates "center" position of frustum. Not weighted by mass.
@@ -90,17 +90,17 @@ class RaySamples(TensorDataclass):
     """Samples along a ray
 
     Args:
-        frustums (Frustums): Frustums along ray.
-        camera_indices (TensorType[..., 1]): Camera index.
-        valid_mask (TensorType[..., 1]): Rays that are valid.
-        deltas (TensorType[..., 1]): "width" of each sample.
-        ts (TensorType[..., 1]): sample-camera distance of each sample.
+        frustums: Frustums along ray.
+        camera_indices: Camera index.
+        valid_mask: Rays that are valid.
+        deltas: "width" of each sample.
+        ts: sample-camera distance of each sample.
     """
 
     frustums: Frustums
-    camera_indices: TensorType[..., 1] = None
-    valid_mask: TensorType[..., 1] = None
-    deltas: TensorType[..., 1] = None
+    camera_indices: TensorType["bs":..., 1] = None
+    valid_mask: TensorType["bs":..., 1] = None
+    deltas: TensorType["bs":..., 1] = None
 
     def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples", 1]:
         """Return weights based on predicted densities
@@ -152,13 +152,13 @@ class RayBundle(TensorDataclass):
     """A bundle of ray parameters.
 
     Args:
-        origins (TensorType[..., 3]): Ray origins (XYZ)
-        directions (TensorType[..., 3]): Unit ray direction vector
-        pixel_area (TensorType[..., 1]): Projected area of pixel a distance 1 away from origin.
-        camera_indices (TensorType[..., 1]): Camera indices
-        nears (TensorType[..., 1]): Distance along ray to start sampling
-        fars (TensorType[..., 1]): Rays Distance along ray to stop sampling
-        valid_mask (TensorType[..., 1]): Rays that are valid
+        origins: Ray origins (XYZ)
+        directions: Unit ray direction vector
+        pixel_area: Projected area of pixel a distance 1 away from origin.
+        camera_indices: Camera indices
+        nears: Distance along ray to start sampling
+        fars: Rays Distance along ray to stop sampling
+        valid_mask: Rays that are valid
     """
 
     origins: TensorType["num_rays", 3]
@@ -167,7 +167,7 @@ class RayBundle(TensorDataclass):
     camera_indices: Optional[TensorType["num_rays", 1]] = None
     nears: Optional[TensorType["num_rays", 1]] = None
     fars: Optional[TensorType["num_rays", 1]] = None
-    valid_mask: Optional[TensorType["num_rays", 1]] = None
+    valid_mask: Optional[TensorType["num_rays", 1, bool]] = None
     num_rays_per_chunk: int = None
 
     def move_to_device(self, device: torch.device) -> None:
@@ -220,7 +220,7 @@ class RayBundle(TensorDataclass):
         return self.flatten()[start_idx:end_idx]
 
     def get_ray_samples(
-        self, bin_starts: TensorType[..., "num_samples", 1], bin_ends: TensorType[..., "num_samples", 1]
+        self, bin_starts: TensorType["bs":..., "num_samples", 1], bin_ends: TensorType["bs":..., "num_samples", 1]
     ) -> RaySamples:
         """Produces samples for each ray by projection points along the ray direction.
 

nerfactory/fields/modules/base.py

@@ -59,12 +59,12 @@ class FieldModule(nn.Module):
         return self.out_dim
 
     @abstractmethod
-    def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
+    def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
         """
         Args:
-            in_tensor (TensorType[..., "input_dim"]): Input tensor to process
+            in_tensor: Input tensor to process
 
         Returns:
-            TensorType[..., "output_dim"]: Processed tensor
+            Processed tensor
         """
         raise NotImplementedError

nerfactory/fields/modules/encoding.py

@@ -33,7 +33,7 @@ class Encoding(FieldModule):
     """Encode an input tensor. Intended to be subclassed
 
     Args:
-        in_dim (int): Input dimension of tensor
+        in_dim: Input dimension of tensor
     """
 
     def __init__(self, in_dim: int) -> None:
@@ -42,7 +42,7 @@ class Encoding(FieldModule):
         super().__init__(in_dim=in_dim)
 
     @abstractmethod
-    def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
+    def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
         raise NotImplementedError
 
 
@@ -52,7 +52,7 @@ class Identity(Encoding):
     def get_out_dim(self) -> int:
         return self.in_dim
 
-    def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
+    def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
         return in_tensor
 
 
@@ -60,9 +60,9 @@ class ScalingAndOffset(Encoding):
     """Simple scaling and offet to input
 
     Args:
-        in_dim (int): Input dimension of tensor
-        scaling (float, optional): Scaling applied to tensor. Defaults to 1.0.
-        offset (float, optional): Offset applied to tensor. Defaults to 0.0.
+        in_dim: Input dimension of tensor
+        scaling: Scaling applied to tensor. Defaults to 1.0.
+        offset: Offset applied to tensor. Defaults to 0.0.
     """
 
     def __init__(self, in_dim: int, scaling: float = 1.0, offset: float = 0.0) -> None:
@@ -74,7 +74,7 @@ class ScalingAndOffset(Encoding):
     def get_out_dim(self) -> int:
         return self.in_dim
 
-    def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
+    def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
         return self.scaling * in_tensor + self.offset
 
 
@@ -83,11 +83,11 @@ class NeRFEncoding(Encoding):
     Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.
 
     Args:
-        in_dim (int): Input dimension of tensor
-        num_frequencies (int): Number of encoded frequencies per axis
-        min_freq_exp (float): Minimum frequency exponent
-        max_freq_exp (float): Maximum frequency exponent
-        include_input (float): Append the input coordinate to the encoding
+        in_dim: Input dimension of tensor
+        num_frequencies: Number of encoded frequencies per axis
+        min_freq_exp: Minimum frequency exponent
+        max_freq_exp: Maximum frequency exponent
+        include_input: Append the input coordinate to the encoding
     """
 
     def __init__(
@@ -107,16 +107,18 @@ class NeRFEncoding(Encoding):
         return out_dim
 
     def forward(
-        self, in_tensor: TensorType[..., "input_dim"], covs: Optional[TensorType[..., "input_dim", "input_dim"]] = None
-    ) -> TensorType[..., "output_dim"]:
+        self,
+        in_tensor: TensorType["bs":..., "input_dim"],
+        covs: Optional[TensorType["bs":..., "input_dim", "input_dim"]] = None,
+    ) -> TensorType["bs":..., "output_dim"]:
         """Calculates NeRF encoding. If covariances are provided the encodings will be integrated as proposed
             in mip-NeRF.
 
         Args:
-            in_tensor (TensorType[..., "input_dim"]): For best performance, the input tensor should be between 0 and 1.
-            covs (TensorType[..., "input_dim", "input_dim"], optional): Covariances of input points. Defaults to None.
+            in_tensor: For best performance, the input tensor should be between 0 and 1.
+            covs: Covariances of input points. Defaults to None.
         Returns:
-            TensorType[..., "output_dim"]: Output values will be between -1 and 1
+            Output values will be between -1 and 1
         """
         in_tensor = 2 * torch.pi * in_tensor  # scale to [0, 2pi]
         freqs = 2 ** torch.linspace(self.min_freq, self.max_freq, self.num_frequencies).to(in_tensor.device)
@@ -141,10 +143,10 @@ class RFFEncoding(Encoding):
     """Random Fourier Feature encoding. Supports integrated encodings.
 
     Args:
-        in_dim (int): Input dimension of tensor
-        num_frequencies (int): Number of encoding frequencies
-        scale (float): Std of Gaussian to sample frequencies. Must be greater than zero
-        include_input (float): Append the input coordinate to the encoding
+        in_dim: Input dimension of tensor
+        num_frequencies: Number of encoding frequencies
+        scale: Std of Gaussian to sample frequencies. Must be greater than zero
+        include_input: Append the input coordinate to the encoding
     """
 
     def __init__(self, in_dim: int, num_frequencies: int, scale: float, include_input: bool = False) -> None:
@@ -162,17 +164,19 @@ class RFFEncoding(Encoding):
         return self.num_frequencies * 2
 
     def forward(
-        self, in_tensor: TensorType[..., "input_dim"], covs: Optional[TensorType[..., "input_dim", "input_dim"]] = None
-    ) -> TensorType[..., "output_dim"]:
+        self,
+        in_tensor: TensorType["bs":..., "input_dim"],
+        covs: Optional[TensorType["bs":..., "input_dim", "input_dim"]] = None,
+    ) -> TensorType["bs":..., "output_dim"]:
         """Calculates RFF encoding. If covariances are provided the encodings will be integrated as proposed
             in mip-NeRF.
 
         Args:
-            in_tensor (TensorType[..., "input_dim"]): For best performance, the input tensor should be between 0 and 1.
-            covs (TensorType[..., "input_dim", "input_dim"], optional): Covariances of input points. Defaults to None.
+            in_tensor: For best performance, the input tensor should be between 0 and 1.
+            covs: Covariances of input points. Defaults to None.
 
         Returns:
-            TensorType[..., "output_dim"]: Output values will be between -1 and 1
+            Output values will be between -1 and 1
         """
         in_tensor = 2 * torch.pi * in_tensor  # scale to [0, 2pi]
         scaled_inputs = in_tensor @ self.b_matrix  # [..., "num_frequencies"]
@@ -195,12 +199,12 @@ class HashEncoding(Encoding):
     """Hash encoding
 
     Args:
-        num_levels (int, optional): Number of feature grids. Defaults to 16.
-        min_res (int, optional): Resolution of smallest feature grid. Defaults to 16.
-        max_res (int, optional): Resolution of largest feature grid. Defaults to 1024.
-        hash_table_size (int, optional): Size of hash table. Defaults to 2**19.
-        features_per_level (int, optional): Number of features per level. Defaults to 2.
-        hash_init_scale (float, optional): Value to initialize hash grid. Defaults to 0.001.
+        num_levels: Number of feature grids. Defaults to 16.
+        min_res: Resolution of smallest feature grid. Defaults to 16.
+        max_res: Resolution of largest feature grid. Defaults to 1024.
+        hash_table_size: Size of hash table. Defaults to 2**19.
+        features_per_level: Number of features per level. Defaults to 2.
+        hash_init_scale: Value to initialize hash grid. Defaults to 0.001.
     """
 
     def __init__(
@@ -230,14 +234,14 @@ class HashEncoding(Encoding):
     def get_out_dim(self) -> int:
         return self.num_levels * self.features_per_level
 
-    def hash_fn(self, in_tensor: TensorType[..., "num_levels", 3]) -> TensorType[..., "num_levels"]:
+    def hash_fn(self, in_tensor: TensorType["bs":..., "num_levels", 3]) -> TensorType["bs":..., "num_levels"]:
         """Hash tensor using method described in Instant-NGP
 
         Args:
-            in_tensor (TensorType[..., "num_levels", 3]): Tensor to be hashed
+            in_tensor: Tensor to be hashed
 
         Returns:
-            TensorType[..., "num_levels"]: Hashed tensor
+            Hashed tensor
         """
 
         # min_val = torch.min(in_tensor)
@@ -252,7 +256,7 @@ class HashEncoding(Encoding):
         x += self.hash_offset.to(x.device)
         return x
 
-    def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
+    def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
         assert in_tensor.shape[-1] == 3
 
         in_tensor = in_tensor[..., None, :]  # [..., 1, 3]
@@ -299,9 +303,9 @@ class TensorCPEncoding(Encoding):
     """Learned CANDECOMP/PARFAC (CP) decomposition encoding used in TensoRF
 
     Args:
-        resolution (int, optional): Resolution of grid. Defaults to 256.
-        num_components (int, optional): Number of components per dimension. Defaults to 24.
-        init_scale (float, optional): Initialization scale. Defaults to 0.1.
+        resolution: Resolution of grid. Defaults to 256.
+        num_components: Number of components per dimension. Defaults to 24.
+        init_scale: Initialization scale. Defaults to 0.1.
     """
 
     def __init__(self, resolution: int = 256, num_components: int = 24, init_scale: float = 0.1) -> None:
@@ -316,7 +320,7 @@ class TensorCPEncoding(Encoding):
     def get_out_dim(self) -> int:
         return self.num_components
 
-    def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
+    def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
         line_coord = torch.stack([in_tensor[..., 2], in_tensor[..., 1], in_tensor[..., 0]])  # [3, ...]
         line_coord = torch.stack([torch.zeros_like(line_coord), line_coord], dim=-1)  # [3, ...., 2]
 
@@ -335,7 +339,7 @@ class TensorCPEncoding(Encoding):
         """Upsamples underyling feature grid
 
         Args:
-            resolution (int): Target resolution.
+            resolution: Target resolution.
         """
 
         self.line_coef.data = F.interpolate(
@@ -349,9 +353,9 @@ class TensorVMEncoding(Encoding):
     """Learned vector-matrix encoding proposed by TensoRF
 
     Args:
-        resolution (int, optional): Resolution of grid. Defaults to 256.
-        num_components (int, optional): Number of components per dimension. Defaults to 24.
-        init_scale (float, optional): Initialization scale. Defaults to 0.1.
+        resolution: Resolution of grid. Defaults to 256.
+        num_components: Number of components per dimension. Defaults to 24.
+        init_scale: Initialization scale. Defaults to 0.1.
     """
 
     def __init__(self, resolution: int = 256, num_components: int = 24, init_scale: float = 0.1) -> None:
@@ -367,7 +371,7 @@ class TensorVMEncoding(Encoding):
     def get_out_dim(self) -> int:
         return self.num_components * 3
 
-    def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
+    def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
         in_tensor = in_tensor / 4.0 + 0.5
         plane_coord = torch.stack([in_tensor[..., [0, 1]], in_tensor[..., [0, 2]], in_tensor[..., [1, 2]]])  # [3,...,2]
         line_coord = torch.stack([in_tensor[..., 2], in_tensor[..., 1], in_tensor[..., 0]])  # [3, ...]
@@ -390,7 +394,7 @@ class TensorVMEncoding(Encoding):
         """Upsamples underyling feature grid
 
         Args:
-            resolution (int): Target resolution.
+            resolution: Target resolution.
         """
 
         self.plane_coef.data = F.interpolate(
@@ -407,7 +411,7 @@ class SHEncoding(Encoding):
     """Spherical harmonic encoding
 
     Args:
-        levels (int, optional): Number of spherical hamonic levels to encode. Defaults to 4.
+        levels: Number of spherical hamonic levels to encode. Defaults to 4.
     """
 
     def __init__(self, levels: int = 4) -> None:
@@ -422,5 +426,5 @@ class SHEncoding(Encoding):
         return self.levels**2
 
     @torch.no_grad()
-    def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
+    def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
         return components_from_spherical_harmonics(levels=self.levels, directions=in_tensor)

nerfactory/fields/modules/field_heads.py

@@ -40,10 +40,10 @@ class FieldHead(FieldModule):
     """Base field output
 
     Args:
-        out_dim (int): output dimension for renderer
-        field_head_name (FieldHeadNames): Field type
-        in_dim (int, optional): input dimension. If not defined in constructor, it must be set later.
-        activation (Optional[nn.Module]): output head activation
+        out_dim: output dimension for renderer
+        field_head_name: Field type
+        in_dim: input dimension. If not defined in constructor, it must be set later.
+        activation: output head activation
     """
 
     def __init__(
@@ -71,14 +71,14 @@ class FieldHead(FieldModule):
     def _construct_net(self):
         self.net = nn.Linear(self.in_dim, self.out_dim)
 
-    def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType:
+    def forward(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":..., "out_dim"]:
         """Process network output for renderer
 
         Args:
-            in_tensor (TensorType[..., "in_dim]): Network input
+            in_tensor: Network input
 
         Returns:
-            TensorType[..., "out_dim"]: Render head output
+            Render head output
         """
         if not self.net:
             raise SystemError("in_dim not set. Must be provided to construtor, or set_in_dim() should be called.")
@@ -92,8 +92,8 @@ class DensityFieldHead(FieldHead):
     """Density output
 
     Args:
-        in_dim (int, optional): input dimension. If not defined in constructor, it must be set later.
-        activation (Optional[nn.Module]): output head activation
+        in_dim: input dimension. If not defined in constructor, it must be set later.
+        activation: output head activation
     """
 
     def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None:
@@ -104,8 +104,8 @@ class RGBFieldHead(FieldHead):
     """RGB output
 
     Args:
-        in_dim (int, optional): input dimension. If not defined in constructor, it must be set later.
-        activation (Optional[nn.Module]): output head activation
+        in_dim: input dimension. If not defined in constructor, it must be set later.
+        activation: output head activation
     """
 
     def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Sigmoid()) -> None:
@@ -116,10 +116,10 @@ class SHFieldHead(FieldHead):
     """Spherical harmonics output
 
     Args:
-        in_dim (int, optional): input dimension. If not defined in constructor, it must be set later.
-        levels (int, optional): Number of spherical harmonics layers. Defaults to 3.
-        channels (int, optional): Number of channels. Defaults to 3 (ie RGB).
-        activation (Optional[nn.Module], optional): Output activation. Defaults to None.
+        in_dim: input dimension. If not defined in constructor, it must be set later.
+        levels: Number of spherical harmonics layers. Defaults to 3.
+        channels: Number of channels. Defaults to 3 (ie RGB).
+        activation: Output activation. Defaults to None.
     """
 
     def __init__(
@@ -134,8 +134,8 @@ class UncertaintyFieldHead(FieldHead):
     """Uncertainty output
 
     Args:
-        in_dim (int, optional): input dimension. If not defined in constructor, it must be set later.
-        activation (Optional[nn.Module]): output head activation
+        in_dim: input dimension. If not defined in constructor, it must be set later.
+        activation: output head activation
     """
 
     def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None:
@@ -146,8 +146,8 @@ class TransientRGBFieldHead(FieldHead):
     """Transient RGB output
 
     Args:
-        in_dim (int, optional): input dimension. If not defined in constructor, it must be set later.
-        activation (Optional[nn.Module]): output head activation
+        in_dim: input dimension. If not defined in constructor, it must be set later.
+        activation: output head activation
     """
 
     def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Sigmoid()) -> None:
@@ -158,8 +158,8 @@ class TransientDensityFieldHead(FieldHead):
     """Transient density output
 
     Args:
-        in_dim (int, optional): input dimension. If not defined in constructor, it must be set later.
-        activation (Optional[nn.Module]): output head activation
+        in_dim: input dimension. If not defined in constructor, it must be set later.
+        activation: output head activation
     """
 
     def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None:
@@ -172,9 +172,9 @@ class SemanticFieldHead(FieldHead):
     """Semantic stuff output
 
     Args:
-        num_classes (int): Number of semantic classes
-        in_dim (int, optional): input dimension. If not defined in constructor, it must be set later.
-        activation (Optional[nn.Module]): output head activation
+        num_classes: Number of semantic classes
+        in_dim: input dimension. If not defined in constructor, it must be set later.
+        activation: output head activation
     """
 
     def __init__(self, num_classes: int, in_dim: Optional[int] = None) -> None:

nerfactory/fields/modules/mlp.py

@@ -28,12 +28,12 @@ class MLP(FieldModule):
     """Multilayer perceptron
 
     Args:
-        in_dim (int): Input layer dimension
-        num_layers (int): Number of network layers
-        layer_width (int): Width of each MLP layer
-        out_dim (int, optional): Ouput layer dimension. Defaults to layer_width
-        activation (Optional[nn.Module], optional): intermediate layer activation function. Defaults to nn.ReLU.
-        out_activation (Optional[nn.Module], optional): output activation function. Defaults to None.
+        in_dim: Input layer dimension
+        num_layers: Number of network layers
+        layer_width: Width of each MLP layer
+        out_dim: Ouput layer dimension. Defaults to layer_width
+        activation: intermediate layer activation function. Defaults to nn.ReLU.
+        out_activation: output activation function. Defaults to None.
     """
 
     def __init__(
@@ -76,14 +76,14 @@ class MLP(FieldModule):
             layers.append(nn.Linear(self.layer_width, self.out_dim))
         self.layers = nn.ModuleList(layers)
 
-    def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]:
+    def forward(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":..., "out_dim"]:
         """Process input with a multilayer perceptron.
 
         Args:
-            in_tensor (TensorType[..., "in_dim]): Network input
+            in_tensor: Network input
 
         Returns:
-            TensorType[..., "out_dim"]: MLP network output
+            MLP network output
         """
         x = in_tensor
         for i, layer in enumerate(self.layers):

nerfactory/fields/modules/spatial_distortions.py

@@ -27,13 +27,15 @@ from nerfactory.utils.math import Gaussians
 class SpatialDistortion(nn.Module):
     """Apply spatial distortions"""
 
-    def forward(self, positions: Union[TensorType[..., 3], Gaussians]) -> Union[TensorType[..., 3], Gaussians]:
+    def forward(
+        self, positions: Union[TensorType["bs":..., 3], Gaussians]
+    ) -> Union[TensorType["bs":..., 3], Gaussians]:
         """
         Args:
-            positions (Union[TensorType[..., 3], Gaussians]): Sample to distort
+            positions: Sample to distort
 
         Returns:
-            Union[TensorType[..., 3], Gaussians]: distorted sample
+            Union: distorted sample
         """
 
 

nerfactory/graphs/base.py

@@ -43,7 +43,7 @@ def setup_graph(config: GraphConfig, dataset_inputs: DatasetInputs, device: str)
     """Setup the graph. The dataset inputs should be set with the training data.
 
     Args:
-        dataset_inputs (DatasetInputs): The inputs which will be used to define the camera parameters.
+        dataset_inputs: The inputs which will be used to define the camera parameters.
     """
     graph = instantiate_from_dict_config(DictConfig(config), **dataset_inputs.as_dict())
     graph.to(device)
@@ -77,14 +77,14 @@ class Graph(AbstractGraph):
     subclassed for custom NeRF model.
 
     Args:
-        intrinsics (torch.Tensor): Camera intrinsics.
-        camera_to_world (torch.Tensor): Camera to world transformation.
-        loss_coefficients (DictConfig): Loss specific weights.
-        scene_bounds (SceneBounds): Bounds of target scene.
-        enable_collider (bool): Whether to create a scene collider to filter rays.
-        collider_config (DictConfig): Configuration of scene collider.
-        enable_density_field (bool): Whether to create a density field to filter samples.
-        density_field_config (DictConfig): Configuration of density field.
+        intrinsics: Camera intrinsics.
+        camera_to_world: Camera to world transformation.
+        loss_coefficients: Loss specific weights.
+        scene_bounds: Bounds of target scene.
+        enable_collider: Whether to create a scene collider to filter rays.
+        collider_config: Configuration of scene collider.
+        enable_density_field: Whether to create a density field to filter samples.
+        density_field_config: Configuration of density field.
     """
 
     def __init__(
@@ -145,7 +145,7 @@ class Graph(AbstractGraph):
         """Obtain the parameter groups for the optimizers
 
         Returns:
-            Dict[str, List[Parameter]]: Mapping of different parameter groups
+            Mapping of different parameter groups
         """
 
     @abstractmethod
@@ -153,10 +153,10 @@ class Graph(AbstractGraph):
         """Takes in a Ray Bundle and returns a dictionary of outputs.
 
         Args:
-            ray_bundle (RayBundle): Input bundle of rays.
+            ray_bundle: Input bundle of rays.
 
         Returns:
-            dict: Outputs of graph. (ie. rendered colors)
+            Outputs of graph. (ie. rendered colors)
         """
 
     def process_outputs_as_images(self, outputs):  # pylint:disable=no-self-use

nerfactory/graphs/modules/ray_generator.py

@@ -27,9 +27,8 @@ class RayGenerator(nn.Module):
     This class will store the intrinsics and extrinsics parameters of the cameras.
 
     Args:
-        intrinsics (TensorType["num_cameras", "num_intrinsics_params"]):
-            The intrinsics parameters.
-        camera_to_world (TensorType["num_cameras", 3, 4]): Camera to world transformation matrix.
+        intrinsics: The intrinsics parameters.
+        camera_to_world: Camera to world transformation matrix.
     """
 
     def __init__(
@@ -54,7 +53,7 @@ class RayGenerator(nn.Module):
         """Index into the cameras to generate the rays.
 
         Args:
-            ray_indices (TensorType["num_rays", 3]): Contains camera, row, and col indicies for target rays.
+            ray_indices: Contains camera, row, and col indicies for target rays.
         """
         c = ray_indices[:, 0]  # camera indices
         y = ray_indices[:, 1]  # row indices

nerfactory/graphs/modules/ray_losses.py

@@ -20,7 +20,9 @@ from torchtyping import TensorType
 from nerfactory.cameras.rays import RaySamples
 
 
-def distortion_loss(ray_samples: RaySamples, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., 1]:
+def distortion_loss(
+    ray_samples: RaySamples, densities: TensorType["bs":..., "num_samples", 1]
+) -> TensorType["bs":..., 1]:
     """Ray baserd distortion loss proposed in MipNeRF-360.
 
     .. math::
@@ -32,11 +34,11 @@ def distortion_loss(ray_samples: RaySamples, densities: TensorType[..., "num_sam
     is the weight at location :math:`u` between bin locations :math:`s_i` and :math:`s_{i+1}`.
 
     Args:
-        ray_samples (RaySamples): Ray samples to compute loss over
-        densities (TensorType[..., "num_samples", 1]): Predicted sample densities
+        ray_samples: Ray samples to compute loss over
+        densities: Predicted sample densities
 
     Returns:
-        TensorType[..., 1]: Distortion Loss.
+        Distortion Loss.
     """
 
     # Compute the weight at each sample location

nerfactory/graphs/modules/ray_sampler.py

@@ -69,13 +69,12 @@ class SpacedSampler(Sampler):
     ) -> None:
         """
         Args:
-            num_samples (int): Number of samples per ray
-            spacing_fn (Callable): Function that dictates sample spacing (ie `lambda x : x` is uniform).
-            spacing_fn_inv (Callable): The inverse of spacing_fn.
-            train_stratified (bool): Use stratified sampling during training. Defults to True
-            density_field (DensityGrid, optional): Density grid. If provides,
-                samples below weight_threshold as set as invalid.
-            weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided.
+            num_samples: Number of samples per ray
+            spacing_fn: Function that dictates sample spacing (ie `lambda x : x` is uniform).
+            spacing_fn_inv: The inverse of spacing_fn.
+            train_stratified: Use stratified sampling during training. Defults to True
+            density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
+            weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
         """
         super().__init__(num_samples=num_samples, density_field=density_field, weight_threshold=weight_threshold)
         self.train_stratified = train_stratified
@@ -91,11 +90,11 @@ class SpacedSampler(Sampler):
         """Generates position samples accoring to spacing function.
 
         Args:
-            ray_bundle (RayBundle): Rays to generate samples for
-            num_samples (Optional[int]): Number of samples per ray
+            ray_bundle: Rays to generate samples for
+            num_samples: Number of samples per ray
 
         Returns:
-            RaySamples: Positions and deltas for samples along a ray
+            Positions and deltas for samples along a ray
         """
         assert ray_bundle is not None
         assert ray_bundle.nears is not None
@@ -133,11 +132,10 @@ class UniformSampler(SpacedSampler):
     ) -> None:
         """
         Args:
-            num_samples (int): Number of samples per ray
-            train_stratified (bool): Use stratified sampling during training. Defults to True
-            density_field (DensityGrid, optional): Density grid. If provides,
-                samples below weight_threshold as set as invalid.
-            weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided.
+            num_samples: Number of samples per ray
+            train_stratified: Use stratified sampling during training. Defults to True
+            density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
+            weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
         """
         super().__init__(
             num_samples=num_samples,
@@ -161,11 +159,10 @@ class LinearDisparitySampler(SpacedSampler):
     ) -> None:
         """
         Args:
-            num_samples (int): Number of samples per ray
-            train_stratified (bool): Use stratified sampling during training. Defults to True
-            density_field (DensityGrid, optional): Density grid. If provides,
-                samples below weight_threshold as set as invalid.
-            weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided.
+            num_samples: Number of samples per ray
+            train_stratified: Use stratified sampling during training. Defults to True
+            density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
+            weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
         """
         super().__init__(
             num_samples=num_samples,
@@ -189,11 +186,10 @@ class SqrtSampler(SpacedSampler):
     ) -> None:
         """
         Args:
-            num_samples (int): Number of samples per ray
-            train_stratified (bool): Use stratified sampling during training. Defults to True
-            density_field (DensityGrid, optional): Density grid. If provides,
-                samples below weight_threshold as set as invalid.
-            weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided.
+            num_samples: Number of samples per ray
+            train_stratified: Use stratified sampling during training. Defults to True
+            density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
+            weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
         """
         super().__init__(
             num_samples=num_samples,
@@ -217,11 +213,10 @@ class LogSampler(SpacedSampler):
     ) -> None:
         """
         Args:
-            num_samples (int): Number of samples per ray
-            train_stratified (bool): Use stratified sampling during training. Defults to True
-            density_field (DensityGrid, optional): Density grid. If provides,
-                samples below weight_threshold as set as invalid.
-            weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided.
+            num_samples: Number of samples per ray
+            train_stratified: Use stratified sampling during training. Defults to True
+            density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
+            weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
         """
         super().__init__(
             num_samples=num_samples,
@@ -247,13 +242,12 @@ class PDFSampler(Sampler):
     ) -> None:
         """
         Args:
-            num_samples (int): Number of samples per ray
-            train_stratified: boolean: Randomize location within each bin during training. Defaults to True
+            num_samples: Number of samples per ray
+            train_stratified: Randomize location within each bin during training. Defaults to True
             include_original: Add original samples to ray. Defaults to True
-            density_field (DensityGrid, optional): Density grid. If provides,
-                samples below weight_threshold as set as invalid.
-            weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided.
-            histogram_padding (float): Amount to weights prior to computing PDF. Defaults to 0.01.
+            density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
+            weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
+            histogram_padding: Amount to weights prior to computing PDF. Defaults to 0.01.
         """
         super().__init__(num_samples=num_samples, density_field=density_field, weight_threshold=weight_threshold)
         self.train_stratified = train_stratified
@@ -272,14 +266,14 @@ class PDFSampler(Sampler):
         """Generates position samples given a distribution.
 
         Args:
-            ray_bundle (RayBundle): Rays to generate samples for
-            ray_samples (RaySamples): Existing ray samples
-            weights: (TensorType[..., "num_samples", 1]): Weights for each bin
-            num_samples (Optional[int]): Number of samples per ray
-            eps: float: Small value to prevent numerical issues. Defaults to 1e-5
+            ray_bundle: Rays to generate samples for
+            ray_samples: Existing ray samples
+            weights: Weights for each bin
+            num_samples: Number of samples per ray
+            eps: Small value to prevent numerical issues. Defaults to 1e-5
 
         Returns:
-            RaySamples: Positions and deltas for samples along a ray
+            Positions and deltas for samples along a ray
         """
         num_samples = num_samples or self.num_samples
         num_bins = num_samples + 1
@@ -367,12 +361,11 @@ class NGPSpacedSampler(Sampler):
         TODO(ruilongli): maybe move aabb test to the collision detector?
 
         Args:
-            ray_bundle (RayBundle): Rays to generate samples for
-            aabb (TensorType[2, 3]): Bounding box of the scene.
-            num_samples (Optional[int]): Number of samples per ray
+            ray_bundle: Rays to generate samples for
+            aabb: Bounding box of the scene.
+            num_samples: Number of samples per ray
 
         Returns:
-            Tuple[RaySamples, TensorType["total_samples", 3], TensorType["total_samples"], TensorType["total_samples"]]:
             First return is ray samples in a packed way where only the valid samples are kept.
             Second return contains all the information to recover packed samples into unpacked mode for rendering.
             The last two returns are t_min and t_max from ray-aabb test.