//  [python]
def interpolate_from_grid_coordinates(samples, grid):
  """Performs trilinear interpolation to estimate the value of a grid function.

  This function makes several assumptions to do the lookup:
  1) The grid is LHW and has evenly spaced samples in the range (0, 1), which
    is really the screen space range [0.5, {L, H, W}-0.5].

  Args:
    samples: Tensor with shape [batch_size, sample_count, 3].
    grid: Tensor with shape [batch_size, length, height, width, 1].

  Returns:
    sample: Tensor with shape [batch_size, sample_count, 1] and type float32.
    mask: Tensor with shape [batch_size, sample_count, 1] and type float32
  """
  batch_size, length, height, width = grid.get_shape().as_list()[:4]
  # These asserts aren't required by the algorithm, but they are currently
  # true for the pipeline:
  assert length == height
  assert length == width
  sample_count = samples.get_shape().as_list()[1]
  tf_util.assert_shape(samples, [batch_size, sample_count, 3],
                       'interpolate_from_grid:samples')
  tf_util.assert_shape(grid, [batch_size, length, height, width, 1],
                       'interpolate_from_grid:grid')
  offset_samples = samples  # Used to subtract 0.5
  lower_coords = tf.cast(tf.math.floor(offset_samples), dtype=tf.int32)
  upper_coords = lower_coords + 1
  alphas = tf.floormod(offset_samples, 1.0)

  maximum_value = grid.get_shape().as_list()[1:4]
  size_per_channel = tf.tile(
      tf.reshape(tf.constant(maximum_value, dtype=tf.int32), [1, 1, 3]),
      [batch_size, sample_count, 1])
  # We only need to check that the floor is at least zero and the ceil is
  # no greater than the max index, because floor round negative numbers to
  # be more negative:
  is_valid = tf.logical_and(lower_coords >= 0, upper_coords < size_per_channel)
  # Validity mask has shape [batch_size, sample_count] and is 1.0 where all of
  # x,y,z are within the [0,1] range of the grid.
  validity_mask = tf.reduce_min(
      tf.cast(is_valid, dtype=tf.float32), axis=2, keep_dims=True)

  lookup_coords = [[[], []], [[], []]]
  corners = [[[], []], [[], []]]
  flattened_grid = tf.reshape(grid, [batch_size, length * height * width])
  for xi, x_coord in enumerate([lower_coords[:, :, 0], upper_coords[:, :, 0]]):
    x_coord = tf.clip_by_value(x_coord, 0, width - 1)
    for yi, y_coord in enumerate([lower_coords[:, :, 1], upper_coords[:, :,
                                                                      1]]):
      y_coord = tf.clip_by_value(y_coord, 0, height - 1)
      for zi, z_coord in enumerate(
          [lower_coords[:, :, 2], upper_coords[:, :, 2]]):
        z_coord = tf.clip_by_value(z_coord, 0, length - 1)
        flat_lookup = z_coord * height * width + y_coord * width + x_coord
        lookup_coords[xi][yi].append(flat_lookup)
        lookup_result = tf.batch_gather(flattened_grid, flat_lookup)
        tf_util.assert_shape(lookup_result, [batch_size, sample_count],
                             'interpolate_from_grid:lookup_result x/8')
        print_op = tf.print('corner xyz=%i, %i, %i' % (xi, yi, zi),
                            lookup_result, '\n', 'flat_lookup:', flat_lookup,
                            '\n\n')
        with tf.control_dependencies([print_op]):
          lookup_result = 1.0 * lookup_result
        corners[xi][yi].append(lookup_result)

  alpha_x, alpha_y, alpha_z = tf.unstack(alphas, axis=2)
  one_minus_alpha_x = 1.0 - alpha_x
  one_minus_alpha_y = 1.0 - alpha_y
  # First interpolate a face along x:
  f00 = corners[0][0][0] * one_minus_alpha_x + corners[1][0][0] * alpha_x
  f01 = corners[0][0][1] * one_minus_alpha_x + corners[1][0][1] * alpha_x
  f10 = corners[0][1][0] * one_minus_alpha_x + corners[1][1][0] * alpha_x
  f11 = corners[0][1][1] * one_minus_alpha_x + corners[1][1][1] * alpha_x
  # Next interpolate a long along y:
  l0 = f00 * one_minus_alpha_y + f10 * alpha_y
  l1 = f01 * one_minus_alpha_y + f11 * alpha_y

  # Finally interpolate a point along z:
  p = l0 * (1.0 - alpha_z) + l1 * alpha_z

  tf_util.assert_shape(p, [batch_size, sample_count], 'interpolate_from_grid:p')

  p = tf.reshape(p, [batch_size, sample_count, 1])
  validity_mask = tf.reshape(validity_mask, [batch_size, sample_count, 1])
  return p, validity_mask