google / google/meridian

Bug: `TypeError` when calling `save_meridian` (Meridian 1.4.0)

Open
#1,364 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.5k
Forks
294
Avg merge
1d 21h
Merged PRs (30d)
47

Description

# Problem
I'm seeing a serialization error when saving a Meridian model with custom priors using:

```python
import numpy as np
import tensorflow_probability as tfp

from schema.serde.meridian_serde import save_meridian
from meridian.model import prior_distribution, spec, model

prior_roi_m = prior_distribution.lognormal_dist_from_mean_std(
np.asarray([...], dtype=np.float32),
np.asarray([...], dtype=np.float32),
)

model_spec = spec.ModelSpec(
prior=prior_distribution.PriorDistribution(roi_m=prior_roi_m)
)

meridian = model.Meridian(input_data=..., model_spec=model_spec)
meridian.sample_prior(1000)

save_meridian(meridian, "./model.binpb")
```
Calling `save_meridian()` on a trained Meridian model raises a `TypeError` originating from `DistributionSerde._to_parameter_value_proto()`.

The serializer encounters a `numpy.ndarray` inside a distribution's parameters and fails with `TypeError: Unsupported type: , [...]`

# Environment
- Meridian version: 1.4.0
- Python: 3.13

# Traceback
```python
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Cell In[30], line 1
----> 1 save_meridian(meridian, "./model.binpb")

File /opt/conda/envs/mmm-se/lib/python3.13/site-packages/schema/serde/meridian_serde.py:357, in save_meridian(mmm, file_path, distribution_function_registry, eda_function_registry)
353 _make_dirs(os.path.dirname(file_path))
355 with _file_open(file_path, 'wb') as f:
356 # Creates an MmmKernel.
--> 357 serialized_kernel = MeridianSerde().serialize(
358 mmm,
359 distribution_function_registry=distribution_function_registry,
360 eda_function_registry=eda_function_registry,
361 )
362 if file_path.endswith('.binpb'):
363 f.write(serialized_kernel.SerializeToString())

File /opt/conda/envs/mmm-se/lib/python3.13/site-packages/schema/serde/meridian_serde.py:114, in MeridianSerde.serialize(self, obj, model_id, meridian_version, include_convergence_info, distribution_function_registry, eda_function_registry)
103 distribution_registry = (
104 distribution_function_registry
105 if distribution_function_registry is not None
106 else function_registry_utils.FunctionRegistry()
107 )
109 eda_function_registry = (
110 eda_function_registry
111 if eda_function_registry is not None
112 else function_registry_utils.FunctionRegistry()
113 )
--> 114 meridian_model_proto = self._make_meridian_model_proto(
115 mmm=obj,
116 model_id=model_id,
117 meridian_version=meridian_version,
118 distribution_function_registry=distribution_registry,
119 eda_function_registry=eda_function_registry,
120 include_convergence_info=include_convergence_info,
121 )
122 any_model = any_pb2.Any()
123 any_model.Pack(meridian_model_proto)

File /opt/conda/envs/mmm-se/lib/python3.13/site-packages/schema/serde/meridian_serde.py:164, in MeridianSerde._make_meridian_model_proto(self, mmm, model_id, meridian_version, distribution_function_registry, eda_function_registry, include_convergence_info)
131 def _make_meridian_model_proto(
132 self,
133 mmm: model.Meridian,
(...) 138 include_convergence_info: bool = False,
139 ) -> meridian_pb.MeridianModel:
140 """Constructs a MeridianModel proto from the TrainedModel.
141
142 Args:
(...) 154 A MeridianModel proto.
155 """
156 model_proto = meridian_pb.MeridianModel(
157 model_id=model_id,
158 model_version=str(meridian_version),
159 hyperparameters=hyperparameters.HyperparametersSerde().serialize(
160 mmm.model_spec
161 ),
162 prior_tfp_distributions=distribution.DistributionSerde(
163 distribution_function_registry
--> 164 ).serialize(mmm.model_spec.prior),
165 inference_data=inference_data.InferenceDataSerde().serialize(
166 mmm.inference_data
167 ),
168 )
169 # For backwards compatibility, only serialize EDA spec if it exists.
170 if hasattr(mmm, 'eda_spec'):

File /opt/conda/envs/mmm-se/lib/python3.13/site-packages/schema/serde/distribution.py:74, in DistributionSerde.serialize(self, obj)
71 if not hasattr(obj, param):
72 continue
73 getattr(proto, param).CopyFrom(
---> 74 self._to_distribution_proto(getattr(obj, param))
75 )
76 proto.function_registry.update(self.function_registry.hashed_registry)
77 return proto

File /opt/conda/envs/mmm-se/lib/python3.13/site-packages/schema/serde/distribution.py:163, in DistributionSerde._to_distribution_proto(self, dist)
158 dist_name = type(dist).__name__
159 dist_class = getattr(backend.tfd, dist_name)
160 return meridian_pb.TfpDistribution(
161 distribution_type=dist_name,
162 parameters={
--> 163 name: self._to_parameter_value_proto(name, value, dist_class)
164 for name, value in dist.parameters.items()
165 },
166 )

File /opt/conda/envs/mmm-se/lib/python3.13/site-packages/schema/serde/distribution.py:262, in DistributionSerde._to_parameter_value_proto(self, param_name, value, dist)
255 raise ValueError(
256 f"Custom function `{param_name}` detected for"
257 f" {type(dist).__name__}, but not found in registry. Please"
258 " add custom functions to registry when saving models."
259 )
261 # Handle unsupported types.
--> 262 raise TypeError(f"Unsupported type: {type(value)}, {value}")

TypeError: Unsupported type: , [ 0.22160271 0.22160271 0.22160271 0.22160271 0.22160271 0.22160271
-0.00124847 0.22160271 0.22160271 0.22160271 0.22160271 0.22160271
0.22160271 0.3527848 0.22160271 0.22160271 0.22160271 0.5815754
0.22160271]
```

Contributor guide

Open the contributing guide

Research direction

Start in schema/serde/distribution.py at DistributionSerde._to_parameter_value_proto(), following the call from schema/serde/meridian_serde.py. Reproduce the failure with the custom prior shown in the issue and inspect how distribution parameters are converted. Done means save_meridian() serializes this numpy.ndarray parameter without raising TypeError.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
backend
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.