import torch


def reflect_z(samples):
  """Reflects the sample locations across the planes specified in xyz.

  Args:
    samples: Tensor with shape [..., 3].

  Returns:
    Tensor with shape [..., 3]. The reflected samples.
  """
  to_keep = samples[..., :-1]
  to_reflect = samples[..., -1:]
  return torch.cat([to_keep, -to_reflect], dim=-1)