`plot_tracks_v2` has bug when plotting with `trackgroup` argument.

I am running [this notebook](https://github.com/google-deepmind/tapnet/blob/main/colabs/tapir_clustering.ipynb) for RoboTAP clustering. After computing the clusters, I am running the following cell:

```py
separation_visibility_trim = clustered['separation_visibility']
separation_tracks_trim = clustered['separation_tracks']

pointtrack_video = viz_utils.plot_tracks_v2(
(demo_videos[demo_episode_ids[0]]).astype(np.uint8),
separation_tracks_trim[demo_episode_ids[0]],
1.0-separation_visibility_trim[demo_episode_ids[0]],
trackgroup=clustered['classes']
)
media.show_video(pointtrack_video, fps=20)
```

However, the plot only shows about 10 points no matter how many points I track, and there are really no clusters to be found. I found that if I comment out `trackgroup`, then the plotting code works correctly and I can see the full range of points (although not colored with cluster ID). I can also verify that clusters are correctly computed by plotting individual frames like so:

```py
separation_visibility_trim = clustered['separation_visibility']
separation_tracks_trim = clustered['separation_tracks']

frame = 35
plt.scatter(
separation_tracks_trim["dummy_id"][:, frame, 0],
separation_tracks_trim["dummy_id"][:, frame, 1],
c=clustered["classes"],
cmap="viridis",
)
plt.imshow(video[frame])
```

It's really only when `trackgroup` is specified that this code does not behave properly. Any ideas of how to fix?