code: def _create_flat_edge_indices(
      flat_cube_indices: torch.Tensor,
      grid_size: Tuple[int, int, int],
  ):
      num_xs = (grid_size[0] - 1) * grid_size[1] * grid_size[2]
      y_offset = num_xs
      num_ys = grid_size[0] * (grid_size[1] - 1) * grid_size[2]
      z_offset = num_xs + num_ys
      return torch.stack(
          [
              # Edges spanning x-axis.
              flat_cube_indices[:, 0] * grid_size[1] * grid_size[2]
              + flat_cube_indices[:, 1] * grid_size[2]
              + flat_cube_indices[:, 2],
              flat_cube_indices[:, 0] * grid_size[1] * grid_size[2]
              + (flat_cube_indices[:, 1] + 1) * grid_size[2]
              + flat_cube_indices[:, 2],
              flat_cube_indices[:, 0] * grid_size[1] * grid_size[2]
              + flat_cube_indices[:, 1] * grid_size[2]
              + flat_cube_indices[:, 2]
              + 1,
              flat_cube_indices[:, 0] * grid_size[1] * grid_size[2]
              + (flat_cube_indices[:, 1] + 1) * grid_size[2]
              + flat_cube_indices[:, 2]
              + 1,
              # Edges spanning y-axis.
              (
                  y_offset
                  + flat_cube_indices[:, 0]
                  * (grid_size[1] - 1)
                  * grid_size[2]
                  + flat_cube_indices[:, 1] * grid_size[2]
                  + flat_cube_indices[:, 2]
              ),
              (
                  y_offset
                  + (flat_cube_indices[:, 0] + 1)
                  * (grid_size[1] - 1)
                  * grid_size[2]
                  + flat_cube_indices[:, 1] * grid_size[2]
                  + flat_cube_indices[:, 2]
              ),
              (
                  y_offset
                  + flat_cube_indices[:, 0]
                  * (grid_size[1] - 1)
                  * grid_size[2]
                  + flat_cube_indices[:, 1] * grid_size[2]
                  + flat_cube_indices[:, 2]
                  + 1
              ),
              (
                  y_offset
                  + (flat_cube_indices[:, 0] + 1)
                  * (grid_size[1] - 1)
                  * grid_size[2]
                  + flat_cube_indices[:, 1] * grid_size[2]
                  + flat_cube_indices[:, 2]
                  + 1
              ),
              # Edges spanning z-axis.
              (
                  z_offset
                  + flat_cube_indices[:, 0]
                  * grid_size[1]
                  * (grid_size[2] - 1)
                  + flat_cube_indices[:, 1] * (grid_size[2] - 1)
                  + flat_cube_indices[:, 2]
              ),
              (
                  z_offset
                  + (flat_cube_indices[:, 0] + 1)
                  * grid_size[1]
                  * (grid_size[2] - 1)
                  + flat_cube_indices[:, 1] * (grid_size[2] - 1)
                  + flat_cube_indices[:, 2]
              ),
              (
                  z_offset
                  + flat_cube_indices[:, 0]
                  * grid_size[1]
                  * (grid_size[2] - 1)
                  + (flat_cube_indices[:, 1] + 1) * (grid_size[2] - 1)
                  + flat_cube_indices[:, 2]
              ),
              (
                  z_offset
                  + (flat_cube_indices[:, 0] + 1)
                  * grid_size[1]
                  * (grid_size[2] - 1)
                  + (flat_cube_indices[:, 1] + 1) * (grid_size[2] - 1)
                  + flat_cube_indices[:, 2]
              ),
          ],
          dim=-1,
      )
  
imports: import math
  from dataclasses import dataclass
  from typing import Dict, Optional, Tuple
  import numpy
  import torch
  import torch.nn.functional
  from torch import nn
  from ...configuration_utils import ConfigMixin, register_to_config
  from ...models import ModelMixin
  from ...utils import BaseOutput
  from .camera import create_pan_cameras
label_desc: Use image node_id 9 for calling a global function with example usage: _create_flat_edge_indices(flat_cube_indices, grid_size) and returns: torch