def test_draw_point2d_with_singleton_color_dims(self, dtype, device):
        """Ensure that plotting behavior is consistent if we have a singleton dim for the color."""
        points = torch.tensor([(1, 3), (2, 4)], device=device)
        # Plot given a color tensor of shape [3]
        color_vec = torch.tensor([5, 10, 15], dtype=torch.float32, device=device)
        vec_img = torch.zeros(3, 8, 8, dtype=torch.uint8, device=device)
        drawn_vec_img = draw_point2d(vec_img, points, color_vec)
        # Plot given a color tensor of shape [3, 1]
        color_mat = torch.unsqueeze(color_vec, dim=1)
        mat_img = vec_img.clone()
        drawn_mat_img = draw_point2d(mat_img, points, color_mat)
        # Ensure that we get the same underlying image back
        assert_close(drawn_vec_img, drawn_mat_img)
Ensure that plotting behavior is consistent if we have a singleton dim for the color.
tests/utils/test_draw.py