//  [python]
def atlas_raymarch_rays_tf(origins, directions, atlas_t, atlas_block_indices_t,
                           atlas_params, scene_params, grid_params):
  """Ray marches rays through a SNeRG scene and returns accumulated RGBA.

  Args:
    origins: A tf.tensor [N, 3] of the ray origins.
    directions: A tf.tensor [N, 3] of ray directions.
    atlas_t: A tensorflow tensor  containing the texture atlas.
    atlas_block_indices_t: A tensorflow tensor containing the indirection grid.
    atlas_params: A dict with params for building and rendering with
      the 3D texture atlas.
    scene_params: A dict for scene specific params (bbox, rotation, resolution).
    grid_params: A dict with parameters describing the high-res voxel grid which
      the atlas is representing.

  Returns:
    rgb: An [N, C] tf.tensor with the colors and features accumulated along each
      ray.
    alpha: An [N, 1] tf.tensor with the alpha value accumuated along each ray.
  """

  # Extract the relevant parameters from the dictionaries.
  worldspace_t_opengl = tf.cast(
      scene_params['worldspace_T_opengl'], dtype=tf.float32)
  min_xyz = tf.cast(scene_params['min_xyz'], dtype=tf.float32)
  voxel_size = tf.cast(grid_params['_voxel_size'], dtype=tf.float32)

  atlas_block_size = atlas_params['atlas_block_size']
  grid_size = tf.cast(grid_params['_grid_size'], dtype=tf.float32)
  data_block_size = atlas_params['_data_block_size']
  num_channels = scene_params['_channels']

  # Set up the rays and transform them to the voxel grid coordinate space.
  num_rays = origins.shape[0]
  opengl_t_worldspace = tf.linalg.inv(worldspace_t_opengl)

  origins_hom = tf.concat([origins, tf.ones_like(origins[Ellipsis, 0:1])], axis=-1)
  origins_hom = tf.reshape(origins_hom, (-1, 4))
  origins_hom = tf.matmul(origins_hom, opengl_t_worldspace)
  origins_opengl = tf.reshape(origins_hom[Ellipsis, 0:3], origins.shape)

  directions_norm = tf.linalg.norm(directions, axis=-1, keepdims=True)
  directions /= directions_norm
  directions_opengl = tf.matmul(directions, opengl_t_worldspace[0:3, 0:3])

  origins_grid = (origins_opengl - min_xyz) / voxel_size
  directions_grid = directions_opengl
  inv_directions_grid = 1.0 / directions_grid

  # Now set the near and far distance of each ray to match the cube which
  # the voxel grid is defined in.
  min_distances, max_distances = rays_aabb_intersection_tf(
      tf.zeros_like(min_xyz), tf.cast(grid_size, dtype=tf.float32),
      origins_grid, inv_directions_grid)

  invalid_mask = min_distances > max_distances
  zero_fill = tf.zeros_like(min_distances)
  min_distances = tf.where(invalid_mask, zero_fill, min_distances)
  max_distances = tf.where(invalid_mask, zero_fill, max_distances)

  # The NeRF near/far bounds have been set for unnormalized ray directions, so
  # we need to scale our bounds here to compensate for normalizing.
  near_in_voxels = directions_norm[Ellipsis, 0] * scene_params['near'] / voxel_size
  far_in_voxels = directions_norm[Ellipsis, 0] * scene_params['far'] / voxel_size
  min_distances = tf.maximum(near_in_voxels, min_distances)
  max_distances = tf.maximum(near_in_voxels, max_distances)
  max_distances = tf.minimum(far_in_voxels, max_distances)

  current_distances = tf.expand_dims(min_distances, -1) + 0.5

  # Finally, set up the accumulation buffers we need for ray marching.
  total_rgb = tf.zeros((num_rays, num_channels), dtype=tf.float32)
  total_visibility = tf.ones((num_rays, 1), dtype=tf.float32)

  init_state = (0, current_distances, total_rgb, total_visibility)
  max_num_steps = tf.math.ceil(tf.linalg.norm(grid_size))

  def raymarch_condition(i, current_distances, _, total_visibility):
    """Proceed until each ray is fully opaque or has left the voxel grid.

    Args:
      i: A integer containing the iteration count.
      current_distances: A [N, 1] tensor containing the t values for each ray.
      _: Dummy parameter for the accumuated RGB color along each ray.
      total_visibility: A [N, 1] tensor with the  accumulated visibility
        (1 - alpha) along each ray.

    Returns:
      False if all of the rays have finished ray marching (exited the scene,
      or saturated alpha). Also returns false if a maximum iteration count
      has been reached (computed as the diagonal of the dense high-res voxel
      grid).
    """
    visibility_mask = total_visibility >= 1.0 / 256.0
    distance_mask = current_distances < tf.expand_dims(max_distances, -1)
    active_mask = tf.math.logical_and(visibility_mask, distance_mask)[Ellipsis, 0]
    return tf.cast(i, tf.float32) < max_num_steps and tf.reduce_max(
        tf.cast(active_mask, dtype=tf.float32)) > 0.0

  def raymarch_body(i, current_distances, total_rgb, total_visibility):
    """Performs a single ray marching step for every ray.

    This is the main body of the ray marching loop factored into a nested
    function, which allows it to be called iteratively inside a TF graph.

    Args:
      i: A integer containing the iteration count.
      current_distances: A [N, 1] tensor containing the t values for each ray.
      total_rgb: Dummy parameter for the accumuated RGB color along each ray.
      total_visibility: A [N, 1] tensor with the  accumulated visibility
        (1 - alpha) along each ray.

    Returns:
      A tuple containing the updated parameters (i, current_distances,
      total_rgb, total_visibility) after the ray marching step.
    """
    positions_grid = origins_grid + tf.math.multiply(directions_grid,
                                                     current_distances)
    positions_atlas_grid = tf.cast(
        tf.floor(positions_grid / data_block_size), dtype=tf.int64)

    # Speedup: Only process rays that are inside the voxel grid.
    epsilon = 0.1
    valid_mask = (positions_grid[Ellipsis, 0] < grid_size[Ellipsis, 0] - 0.5 - epsilon)
    valid_mask = tf.math.logical_and(
        valid_mask, positions_grid[Ellipsis, 1] < grid_size[Ellipsis, 1] - 0.5 - epsilon)
    valid_mask = tf.math.logical_and(
        valid_mask, positions_grid[Ellipsis, 2] < grid_size[Ellipsis, 2] - 0.5 - epsilon)
    valid_mask = tf.math.logical_and(valid_mask,
                                     positions_grid[Ellipsis, 0] > 0.5 + epsilon)
    valid_mask = tf.math.logical_and(valid_mask,
                                     positions_grid[Ellipsis, 1] > 0.5 + epsilon)
    valid_mask = tf.math.logical_and(valid_mask,
                                     positions_grid[Ellipsis, 2] > 0.5 + epsilon)
    invalid_mask = tf.math.logical_not(valid_mask)

    # Fetch the atlas indices from the indirection grid.
    positions_atlas_grid = tf.where(
        tf.expand_dims(invalid_mask, -1), tf.zeros_like(positions_atlas_grid),
        positions_atlas_grid)
    block_indices = tf.gather_nd(atlas_block_indices_t, positions_atlas_grid)
    empty_atlas_mask = block_indices[Ellipsis, 0] < 0

    # Compute where each ray intersects the current macroblock.
    min_aabb_positions = tf.cast(
        positions_atlas_grid * data_block_size, dtype=tf.float32)
    max_aabb_positions = min_aabb_positions + data_block_size
    _, max_distance_to_aabb = rays_aabb_intersection_tf(
        min_aabb_positions, max_aabb_positions, origins_grid,
        inv_directions_grid)

    # And then skip past empty macroblocks.
    skip_ahead_mask = tf.math.logical_and(empty_atlas_mask, valid_mask)
    skip_ahead_delta = tf.expand_dims(max_distance_to_aabb,
                                      -1) - current_distances
    skip_ahead_delta = tf.where(
        tf.expand_dims(tf.logical_not(skip_ahead_mask), -1),
        tf.zeros_like(skip_ahead_delta), skip_ahead_delta)

    current_distances += skip_ahead_delta
    current_distances += 1.0

    # Early out if all rays are outside the voxel grid.
    if tf.reduce_max(tf.cast(valid_mask, dtype=tf.float32)) == 0.0:
      return (i + 1, current_distances, total_rgb, total_visibility)

    # For the rays that are 1) inside the voxel grid and 2) inside a non-empty
    # macroblock, we fetch RGB, Features and alpha from the texture atlas.
    block_indices = tf.where(
        tf.expand_dims(empty_atlas_mask, -1), tf.zeros_like(block_indices),
        block_indices)
    block_indices = tf.where(
        tf.expand_dims(invalid_mask, -1), tf.zeros_like(block_indices),
        block_indices)

    positions_atlas = positions_grid - min_aabb_positions
    positions_atlas += tf.cast(
        block_indices * atlas_block_size, dtype=tf.float32)
    positions_atlas += 1.0  # Account for the one-voxel padding in the atlas.

    positions_atlas -= 0.5
    rgb = tf.zeros((num_rays, num_channels), dtype=tf.float32)
    alpha = tf.zeros((num_rays, 1), dtype=tf.float32)

    # Use trilinear interpolation for smoother results.
    offsets_xyz = [(x, y, z)  # pylint: disable=g-complex-comprehension
                   for x in [0.0, 1.0]
                   for y in [0.0, 1.0]
                   for z in [0.0, 1.0]]
    for dx, dy, dz in offsets_xyz:
      offset_t = tf.convert_to_tensor([dx, dy, dz], dtype=tf.float32)
      atlas_indices = tf.cast(
          tf.floor(positions_atlas + offset_t), dtype=tf.int64)
      weights = 1 - tf.abs(
          tf.cast(atlas_indices, dtype=tf.float32) - positions_atlas)
      weights = tf.reshape(tf.math.reduce_prod(weights, axis=-1), alpha.shape)

      atlas_indices = tf.where(
          tf.expand_dims(invalid_mask, -1), tf.zeros_like(atlas_indices),
          atlas_indices)

      gathered_features = tf.cast(
          tf.gather_nd(atlas_t, atlas_indices), tf.float32)
      trilinear_features = gathered_features[Ellipsis, 0:num_channels]
      trilinear_alpha = gathered_features[Ellipsis, num_channels:num_channels + 1]

      rgb += weights * tf.reshape(trilinear_features, rgb.shape)
      alpha += weights * tf.reshape(trilinear_alpha, alpha.shape)

    # Finally, pad use dummy values for rays that were outside the voxel grid
    # or inside empty macroblocks.
    alpha = tf.where(
        tf.expand_dims(invalid_mask, -1), tf.zeros_like(alpha), alpha)
    alpha = tf.where(
        tf.expand_dims(empty_atlas_mask, -1), tf.zeros_like(alpha), alpha)
    rgb = tf.where(tf.expand_dims(invalid_mask, -1), tf.zeros_like(rgb), rgb)
    rgb = tf.where(
        tf.expand_dims(empty_atlas_mask, -1), tf.zeros_like(rgb), rgb)

    # Accumulate RGB, features and alpha, and loop again.
    total_rgb += total_visibility * rgb
    total_visibility *= 1.0 - alpha

    return (i + 1, current_distances, total_rgb, total_visibility)

  _, current_distances, total_rgb, total_visibility = tf.while_loop(
      raymarch_condition, raymarch_body, init_state, parallel_iterations=1)

  return total_rgb, 1.0 - total_visibility