allow tensor in DiagonalNormal dimension

[arrjon]

This pull request updates the DiagonalNormal distribution implementation to support multi-dimensional (non-vector) distributions instead of being limited to vectors. The main changes involve replacing the single dim attribute with a tuple dims, and updating all relevant logic to handle arbitrary shapes.

This is necessary if we pass not a vector but something else to inference networks like diffusion, flow matching etc.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
bayesflow/approximators/continuous_approximator.py	91.87% <100.00%> (+0.12%)	⬆️
bayesflow/distributions/diagonal_normal.py	95.34% <100.00%> (ø)

[stefanradev93]

Hi Jonas, thanks for the PR. See individual comments.

[arrjon]

I changed the batch shape now to

target_dim = self.inference_network.base_distribution.dims batch_shape = keras.ops.shape(inference_conditions)[: -len(target_dim)] 

that should cover now more use cases as the shape is inferred from the conditions, but the base shape is also considered.