text: from typing import Tuple
  
  from tqdm import tqdm
  import numpy as np
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  
  def cast_rays(
      height: int,
      width: int,
      focal: float,
      pose: torch.Tensor,
      device: torch.device = "cpu",
  ):
      ii, jj = torch.meshgrid(
          torch.arange(width, device=device),
          torch.arange(height, device=device),
          indexing="xy",
      )
      directions = torch.stack(
          [
              (ii - width * 0.5) / focal,
              -(jj - height * 0.5) / focal,
              -torch.ones_like(ii),
          ],
          dim=-1,
      )
      directions = torch.sum(directions[..., None, :] * pose[:3, :3], dim=-1)
      origins = pose[:3, -1].expand(directions.shape)
      return origins, directions
  
  
  def sample_points(
      origins: torch.Tensor,
      directions: torch.Tensor,
      z_near: float,
      z_far: float,
      num_samples: int,
      random: bool = True,
      device: torch.device = "cpu",
  ):
      z = torch.linspace(z_near, z_far, num_samples, device=device)
      if random:
          noise = torch.rand(
              origins.shape[0], origins.shape[1], num_samples, device=device
          )
          z = z + noise * (z_far - z_near) / num_samples
      points = origins[..., None, :] + directions[..., None, :] * z[..., :, None]
      return points, z
  
  
  def positional_encoding(
      x: torch.Tensor, encoding_size: int = 6, include_input: bool = True
  ):
      encoding = [x] if include_input else []
      for i in range(encoding_size):
          for fn in [torch.sin, torch.cos]:
              encoding.append(fn((2**i) * x))
      return torch.cat(encoding, dim=-1)
  
  
  def render_volume(
      samples: torch.Tensor, z: torch.Tensor
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      rgb = torch.sigmoid(samples[..., :3])
      sigma_a = F.relu(samples[..., 3])
      dist = z[..., 1:] - z[..., :-1]
      dist = torch.cat([dist, torch.full_like(z[..., :1], 1e10)], dim=-1)
      alpha = 1.0 - torch.exp(-sigma_a * dist)
      weights = alpha * torch.cumprod(1.0 - alpha + 1e-10, dim=-1)
      rgb = (weights[..., None] * rgb).sum(dim=-2)
      return rgb
  
  
  class NeRF(nn.Module):
      def __init__(self, encoding_size: int = 6, hidden_channels: int = 256):
          super().__init__()
          in_channels = 3 + 3 * 2 * encoding_size
          self.fc1 = nn.Linear(in_channels, hidden_channels)
          self.fc2 = nn.Linear(hidden_channels, hidden_channels)
          self.fc3 = nn.Linear(hidden_channels, hidden_channels)
          self.fc4 = nn.Linear(hidden_channels, hidden_channels // 2)
          self.fc5 = nn.Linear(hidden_channels // 2, hidden_channels // 2)
          self.alpha = torch.nn.Linear(hidden_channels, 1)
          self.rgb = torch.nn.Linear(hidden_channels // 2, 3)
  
      def forward(self, x):
          x = F.relu(self.fc1(x))
          x = F.relu(self.fc2(x))
          x = self.fc3(x)
          alpha = self.alpha(x)
          x = F.relu(self.fc4(x))
          x = F.relu(self.fc5(x))
          rgb = self.rgb(x)
          return torch.cat([rgb, alpha], dim=1)
  
  
  if __name__ == "__main__":
      device = "cuda"
      num_iters = 20000
      z_near = 2.0
      z_far = 6.0
      num_samples = 64  # Number of depth samples per ray
      chunk_size = 4096
      encoding_size = 10
      lr = 5e-3
      eval_freq = 100
  
      data = np.load("tinynerf.npz")
      images = torch.from_numpy(data["images"])
      poses = torch.from_numpy(data["poses"])
      focal = torch.tensor(data["focal"])
  
      height, width = images.shape[1:3]
      dataset_size = images.shape[0]
  
      model = NeRF(encoding_size).to(device)
      optimizer = torch.optim.Adam(model.parameters(), lr=lr)
  
      pbar = tqdm(total=num_iters)
      for i in range(num_iters):
          choice = np.random.randint(100)
          image = images[choice].to(device)
          pose = poses[choice].to(device)
          origins, directions = cast_rays(height, width, focal, pose, device=device)
          points, z = sample_points(
              origins, directions, z_near, z_far, num_samples, device=device
          )
          points = points.reshape(-1, 3)
          encoding = positional_encoding(points, encoding_size)
          batches = torch.split(encoding, chunk_size)
          samples = [model(batch) for batch in batches]
          samples = torch.cat(samples, dim=0)
          samples = samples.reshape(height, width, num_samples, 4)
          rgb = render_volume(samples, z)
  
          optimizer.zero_grad()
          loss = F.mse_loss(rgb, image)
          loss.backward()
          optimizer.step()
  
          pbar.set_postfix(loss=loss.item())
          pbar.update()
      pbar.close()
  