content: from dataclasses import dataclass
  from typing import Optional
  
  import torch
  import torch.nn as nn
  from einops import rearrange
  
  from ..utils import BaseModule
  
  
  class TriplaneUpsampleNetwork(BaseModule):
      @dataclass
      class Config(BaseModule.Config):
          in_channels: int
          out_channels: int
  
      cfg: Config
  
      def configure(self) -> None:
          self.upsample = nn.ConvTranspose2d(
              self.cfg.in_channels, self.cfg.out_channels, kernel_size=2, stride=2
          )
  
      def forward(self, triplanes: torch.Tensor) -> torch.Tensor:
          triplanes_up = rearrange(
              self.upsample(
                  rearrange(triplanes, "B Np Ci Hp Wp -> (B Np) Ci Hp Wp", Np=3)
              ),
              "(B Np) Co Hp Wp -> B Np Co Hp Wp",
              Np=3,
          )
          return triplanes_up
  
  
  class NeRFMLP(BaseModule):
      @dataclass
      class Config(BaseModule.Config):
          in_channels: int
          n_neurons: int
          n_hidden_layers: int
          activation: str = "relu"
          bias: bool = True
          weight_init: Optional[str] = "kaiming_uniform"
          bias_init: Optional[str] = None
  
      cfg: Config
  
      def configure(self) -> None:
          layers = [
              self.make_linear(
                  self.cfg.in_channels,
                  self.cfg.n_neurons,
                  bias=self.cfg.bias,
                  weight_init=self.cfg.weight_init,
                  bias_init=self.cfg.bias_init,
              ),
              self.make_activation(self.cfg.activation),
          ]
          for i in range(self.cfg.n_hidden_layers - 1):
              layers += [
                  self.make_linear(
                      self.cfg.n_neurons,
                      self.cfg.n_neurons,
                      bias=self.cfg.bias,
                      weight_init=self.cfg.weight_init,
                      bias_init=self.cfg.bias_init,
                  ),
                  self.make_activation(self.cfg.activation),
              ]
          layers += [
              self.make_linear(
                  self.cfg.n_neurons,
                  4,  # density 1 + features 3
                  bias=self.cfg.bias,
                  weight_init=self.cfg.weight_init,
                  bias_init=self.cfg.bias_init,
              )
          ]
          self.layers = nn.Sequential(*layers)
  
      def make_linear(
          self,
          dim_in,
          dim_out,
          bias=True,
          weight_init=None,
          bias_init=None,
      ):
          layer = nn.Linear(dim_in, dim_out, bias=bias)
  
          if weight_init is None:
              pass
          elif weight_init == "kaiming_uniform":
              torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu")
          else:
              raise NotImplementedError
  
          if bias:
              if bias_init is None:
                  pass
              elif bias_init == "zero":
                  torch.nn.init.zeros_(layer.bias)
              else:
                  raise NotImplementedError
  
          return layer
  
      def make_activation(self, activation):
          if activation == "relu":
              return nn.ReLU(inplace=True)
          elif activation == "silu":
              return nn.SiLU(inplace=True)
          else:
              raise NotImplementedError
  
      def forward(self, x):
          inp_shape = x.shape[:-1]
          x = x.reshape(-1, x.shape[-1])
  
          features = self.layers(x)
          features = features.reshape(*inp_shape, -1)
          out = {"density": features[..., 0:1], "features": features[..., 1:4]}
  
          return out
  