Closed
Description
If I am not mistaken, the default flow matching configuration on dev
seems to be broken currently. This probably also concerns the 2.0.0
release.
Running the two moons notebook gives the following error message:
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Cell In[12], line 1
----> 1 history = flow_matching_workflow.fit_offline(
2 training_data,
3 epochs=epochs,
4 batch_size=batch_size,
5 validation_data=validation_data
6 )
File [~/Programming/IWR/bf2/bayesflow/workflows/basic_workflow.py:714](http://127.0.0.1:8892/lab/tree/examples/bayesflow/workflows/basic_workflow.py#line=713), in BasicWorkflow.fit_offline(self, data, epochs, batch_size, keep_optimizer, validation_data, **kwargs)
679 """
680 Train the approximator offline using a fixed dataset. This approach will be faster than online training,
681 since no computation time is spent in generating new data for each batch, but it assumes that simulations
(...)
709 metric evolution over epochs.
710 """
712 dataset = OfflineDataset(data=data, batch_size=batch_size, adapter=self.adapter)
--> 714 return self._fit(
715 dataset, epochs, strategy="online", keep_optimizer=keep_optimizer, validation_data=validation_data, **kwargs
716 )
File [~/Programming/IWR/bf2/bayesflow/workflows/basic_workflow.py:913](http://127.0.0.1:8892/lab/tree/examples/bayesflow/workflows/basic_workflow.py#line=912), in BasicWorkflow._fit(self, dataset, epochs, strategy, keep_optimizer, validation_data, **kwargs)
910 self.approximator.compile(optimizer=self.optimizer, metrics=kwargs.pop("metrics", None))
912 try:
--> 913 self.history = self.approximator.fit(
914 dataset=dataset, epochs=epochs, validation_data=validation_data, **kwargs
915 )
916 self._on_training_finished()
917 return self.history
File [~/Programming/IWR/bf2/bayesflow/approximators/continuous_approximator.py:200](http://127.0.0.1:8892/lab/tree/examples/bayesflow/approximators/continuous_approximator.py#line=199), in ContinuousApproximator.fit(self, *args, **kwargs)
148 def fit(self, *args, **kwargs):
149 """
150 Trains the approximator on the provided dataset or on-demand data generated from the given simulator.
151 If `dataset` is not provided, a dataset is built from the `simulator`.
(...)
198 If both `dataset` and `simulator` are provided or neither is provided.
199 """
--> 200 return super().fit(*args, **kwargs, adapter=self.adapter)
File [~/Programming/IWR/bf2/bayesflow/approximators/approximator.py:137](http://127.0.0.1:8892/lab/tree/examples/bayesflow/approximators/approximator.py#line=136), in Approximator.fit(self, dataset, simulator, **kwargs)
135 mock_data = dataset[0]
136 mock_data = keras.tree.map_structure(keras.ops.convert_to_tensor, mock_data)
--> 137 self.build_from_data(mock_data)
139 return super().fit(dataset=dataset, **kwargs)
File [~/Programming/IWR/bf2/bayesflow/approximators/approximator.py:26](http://127.0.0.1:8892/lab/tree/examples/bayesflow/approximators/approximator.py#line=25), in Approximator.build_from_data(self, data)
25 def build_from_data(self, data: Mapping[str, any]) -> None:
---> 26 self.compute_metrics(**data, stage="training")
27 self.built = True
File [~/Programming/IWR/bf2/bayesflow/approximators/continuous_approximator.py:135](http://127.0.0.1:8892/lab/tree/examples/bayesflow/approximators/continuous_approximator.py#line=134), in ContinuousApproximator.compute_metrics(self, inference_variables, inference_conditions, summary_variables, sample_weight, stage)
133 # Force a conversion to Tensor
134 inference_variables = keras.tree.map_structure(keras.ops.convert_to_tensor, inference_variables)
--> 135 inference_metrics = self.inference_network.compute_metrics(
136 inference_variables, conditions=inference_conditions, sample_weight=sample_weight, stage=stage
137 )
139 loss = inference_metrics.get("loss", keras.ops.zeros(())) + summary_metrics.get("loss", keras.ops.zeros(()))
141 inference_metrics = {f"{key}[/inference_](http://127.0.0.1:8892/inference_){key}": value for key, value in inference_metrics.items()}
File [~/Programming/IWR/bf2/bayesflow/networks/flow_matching/flow_matching.py:263](http://127.0.0.1:8892/lab/tree/examples/bayesflow/networks/flow_matching/flow_matching.py#line=262), in FlowMatching.compute_metrics(self, x, conditions, sample_weight, stage)
256 x0 = self.base_distribution.sample(keras.ops.shape(x1)[:-1])
258 if self.use_optimal_transport:
259 # we must choose between resampling x0 or x1
260 # since the data is possibly noisy and may contain outliers, it is better
261 # to possibly drop some samples from x1 than from x0
262 # in the marginal over multiple batches, this is not a problem
--> 263 x0, x1, assignments = optimal_transport(
264 x0,
265 x1,
266 seed=self.seed_generator,
267 **self.optimal_transport_kwargs,
268 return_assignments=True,
269 )
270 if conditions is not None:
271 # conditions must be resampled along with x1
272 conditions = keras.ops.take(conditions, assignments, axis=0)
File [~/Programming/IWR/bf2/bayesflow/utils/optimal_transport/optimal_transport.py:41](http://127.0.0.1:8892/lab/tree/examples/bayesflow/utils/optimal_transport/optimal_transport.py#line=40), in optimal_transport(x1, x2, method, return_assignments, **kwargs)
14 def optimal_transport(x1, x2, method="log_sinkhorn", return_assignments=False, **kwargs):
15 """Matches elements from x2 onto x1, such that the transport cost between them is minimized, according to the method
16 and cost matrix used.
17
(...)
39 x1 and x2 in optimal transport permutation order.
40 """
---> 41 assignments = methods[method.lower()](x1, x2, **kwargs)
42 x2 = keras.ops.take(x2, assignments, axis=0)
44 if return_assignments:
File [~/Programming/IWR/bf2/bayesflow/utils/optimal_transport/sinkhorn.py:35](http://127.0.0.1:8892/lab/tree/examples/bayesflow/utils/optimal_transport/sinkhorn.py#line=34), in sinkhorn(x1, x2, seed, **kwargs)
11 def sinkhorn(x1: Tensor, x2: Tensor, seed: int = None, **kwargs) -> (Tensor, Tensor):
12 """
13 Matches elements from x2 onto x1 using the Sinkhorn-Knopp algorithm.
14
(...)
33
34 """
---> 35 plan = sinkhorn_plan(x1, x2, **kwargs)
36 assignments = keras.random.categorical(plan, num_samples=1, seed=seed)
37 assignments = keras.ops.squeeze(assignments, axis=1)
TypeError: sinkhorn_plan() got an unexpected keyword argument 'cost'