Skip to content

Error running Two Moons notebook: flow matching default config produces error #429

Closed
@vpratz

Description

@vpratz

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'

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions