`GaussianProcess` sample does not work with new index points

A minimal example is the following:

import numpy as np
import tensorflow_probability as tfp
tfd = tfp.distributions
tfk = tfp.math.psd_kernels
index_points = np.array([1., 2, 3])[:, None]
other_index_points = np.array([1., 2])[:, None]

gp = tfd.GaussianProcess(tfk.ExponentiatedQuadratic(), index_points=index_points)
# gp.sample(index_points=other_index_points)  # raise ValueError: Tensor's shape (2,) is not compatible with supplied shape (3,)

The GaussianProcess class does implement _sample_n to take into account new index points:

  def _sample_n(self, n, seed=None, index_points=None):
    return self.get_marginal_distribution(index_points).sample(n, seed=seed)

Following the chain of calls, sample calls _call_sample_n, which calls _sample_n (successful) and then _set_sample_static_shape (in distribution.py -- call is unsuccessful, because it calls self.event_shape -- hence the information on the new index points is lost), which fails at setting the shape of the sample.

A workaround is to use

gp.get_marginal_distribution(index_points=other_index_points).sample()  # works fine

and that it's what I have been doing until now.
However, I wanted to extend the GaussianProcess class and override only the _sample_n method instead of the general sample to retain the various checks, but I could not find a way around (instead of reimplementing sample and avoiding the checks).

Needless to say, if we define

other_index_points = np.array([1., 2, 4])[:, None]

then everything works fine (the index points are different, but the event_shape is the same).

Is there something I missed?

Reproduce the example using GaussianProcess.sample with a different number of index points. Start in GaussianProcess._sample_n and distribution.py at _set_sample_static_shape, then trace how event_shape is used. Done means sampling with new index points succeeds when their shape differs, while the existing shape checks remain intact.