AttributeError: module 'neural_diffusion_processes.sde' has no attribute 'conditional_sample2'

[adam-hartshorne]

When attempting to run gaussian-process-interpolation example, I get the following error. It appears that conditional_sample2 is defined in sde_with_mask.py not sde.py

Traceback (most recent call last):
  File "/media/adam/shared_drive/PycharmProjects/neural_diffusion_processes-main/notebooks/gaussian-process-interpolation.py", line 336, in <module>
    ys, ys_steered = steer(I), steer(R)
  File "/media/adam/shared_drive/PycharmProjects/neural_diffusion_processes-main/notebooks/gaussian-process-interpolation.py", line 332, in steer
    return jax.jit(
  File "/media/adam/shared_drive/PycharmProjects/neural_diffusion_processes-main/notebooks/gaussian-process-interpolation.py", line 333, in <lambda>
    jax.vmap(lambda key: conditional_sample(x_known - shift, y_known, x_test - shift, key=key))
  File "/media/adam/shared_drive/PycharmProjects/neural_diffusion_processes-main/notebooks/gaussian-process-interpolation.py", line 309, in <lambda>
    lambda x, y, x_eval, key: ndp.sde.conditional_sample2(
AttributeError: module 'neural_diffusion_processes.sde' has no attribute 'conditional_sample2'

If you do import this function correctly, it results in the following error

TypeError: conditional_sample2() missing 2 required keyword-only arguments: 'mask_context' and 'mask_test'