ValueError: The filename must end in .weights.h5 after the updating tensorflow-probability[tf]

[enigne]

Hi,
I'm using deepxde in my package and recently had issue with the latest updated version of tfp: https://github.com/enigne/PINNICLE/actions/runs/8476374729/job/23225809520

Even if I explicitly use tensorflow-probability[tf] as suggested here , I still got the error when using deepxde.model.Model.save

tests/test_pinn.py:109: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
PINNICLE/pinn.py:164: in train
    self.save_model()
PINNICLE/pinn.py:134: in save_model
    self.model.save(f"{path}/{subfolder}/{name}")
/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/deepxde/model.py:1041: in save
    self.net.save_weights(save_path)
/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/keras/src/utils/traceback_utils.py:122: in error_handler
    raise e.with_traceback(filtered_tb) from None
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <FNN name=fnn_9, built=True>
filepath = '/tmp/pytest-of-runner/pytest-0/test_train0/pinn/model-10.ckpt'
overwrite = True

    @traceback_utils.filter_traceback
    def save_weights(self, filepath, overwrite=True):
        """Saves all layer weights to a `.weights.h5` file.
    
        Args:
            filepath: `str` or `pathlib.Path` object.
                Path where to save the model. Must end in `.weights.h5`.
            overwrite: Whether we should overwrite any existing model
                at the target location, or instead ask the user
                via an interactive prompt.
        """
        if not str(filepath).endswith(".weights.h5"):
>           raise ValueError(
                "The filename must end in `.weights.h5`. "
                f"Received: filepath={filepath}"
            )
E           ValueError: The filename must end in `.weights.h5`. Received: filepath=/tmp/pytest-of-runner/pytest-0/test_train0/pinn/model-10.ckpt

[lululxvi]

The new tensorflow-probability may requrie this h5 format. I haven't tried the new version yet. How about using the older version of TF/tensorflow-probability? Or would you like to send a PR to update the code?