* fix trainable parameters in distributions (#520)

* Improve numerical precision in MVNScore.log_prob

* add log_gamma diagnostic (#522)

* add log_gamma diagnostic

* add missing export for log_gamma

* add missing export for gamma_null_distribution, gamma_discrepancy

* fix broken unit tests

* rename log_gamma module to sbc

* add test_log_gamma unit test

* add return information to log_gamma doc string

* fix typo in docstring, use fixed-length np array to collect log_gammas instead of appending to an empty list

* Breaking changes: Fix bugs regarding counts in standardization layer (#525)

* standardization: add test for multi-input values (failing)

This test reveals to bugs in the standarization layer:

- count is updated multiple times
- batch_count is too small, as the sizes from reduce_axes have to be
  multiplied

* breaking: fix bugs regarding count in standardization layer

Fixes #524

This fixes the two bugs described in c4cc133:

- count was accidentally updated, leading to wrong values
- count was calculated wrongly, as only the batch size was used. Correct
  is the product of all reduce dimensions. This lead to wrong standard
  deviations

While the batch dimension is the same for all inputs, the size of the
second dimension might vary. For this reason, we need to introduce an
input-specific `count` variable. This breaks serialization.

* fix assert statement in test

* rename log_gamma to calibration_log_gamma (#527)

* simple fix

* Hotfix: numercial stability of non-log-stabilized sinkhorn plan (#531)

* fix numerical stability issues in sinkhorn plan

* improve test suite

* fix ultra-strict convergence criterion in log_sinkhorn_plan

* update dependencies

* add comment about convergence check

* update docsting to reflect fixes

* sinkhorn_plan now returns a transport plan with uniform marginal distributions

* add unit test for sinkhorn_plan

* fix sinkhorn function by sampling from the logits of the transpose of the plan, instead of the plan directly

* sinkhorn(x1, x2) now samples from log(plan) to receive assignments such that x2[assignments] matches x1

* re-enable test_assignment_is_optimal() for method='sinkhorn'

* log_sinkhorn now correctly uses log_plan instead of keras.ops.exp(log_plan), log_sinkhorn_plan returns logits of the transport plan

* add unit tests for log_sinkhorn_plan

* fix faulty indexing with tensor for tensorflow backend

* re-add numItermax for ot pot test

---------

Co-authored-by: Daniel Habermann <[email protected]>

* isinstance sequence

* Pass correct training stage in compute_metrics (#534)

* Pass correct training stage in CouplingFlow.compute_metrics

* Pass correct training stage in CIF and PointInferenceNetwork

* Custom test quantity support for calibration_ecdf (#528)

* Custom test quantity support for calibration_ecdf

* rename variable [no ci]

* Consistent defaults for variable_keys/names in calibration_ecdf with test quantiles

* Tests for calibration_ecdf with test_quantities

* Remove redundant and simplify comments

* Fix docstrings and typehints

---------

Co-authored-by: stefanradev93 <[email protected]>

* Log gamma test fix (#535)

* fix test_calibration_log_gamma_end_to_end unit test failing too often than expected

* set alpha to 0.1% in binom.ppf

* fix typo in comment

* Stateless adapters (#536)

* Remove stateful adapter features

* Fix tests

* Fix typo

* Remove nnpe from adapter

* Bring back notes [skip ci]

* Remove unncessary restriction to kwargs only [skip ci]

* Remove old super call [skip ci]

* Robustify type [skip ci]

* remove standardize from multimodal sim notebook [no ci]

* add draft module docstring to augmentations module [no ci]

Feel free to modify.

* adapt and run neurocognitive modeling notebook [no ci]

* adapt cCM playground notebook [no ci]

* adapt signature of Adapter.standardize

* add parameters missed in previous commit

* Minor NNPE polishing

* remove stage in docstring from OnlineDataset

---------

Co-authored-by: Lasse Elsemüller <[email protected]>
Co-authored-by: Valentin Pratz <[email protected]>

* Fix training strategies in BasicWorkflow

* move multimodal data notebook to regular examples [no ci]

* make pip install call on homepage more verbose [no ci]

* remove deprecated summaries function

The function was renamed to summarize in v2.0.4.

* detail subsampling behavior docs for SIR simulator [no ci]

fixes #518

* move DiffusionModel from experimental to networks

Stabilizes the DiffusionModel class. A deprecation warning for the
DiffusionModel class in the experimental module was added.

* Add citation for resnet (#537) [no ci]

* added citation for resnet

* minor formatting

---------

Co-authored-by: Valentin Pratz <[email protected]>

* Bump up version [skip ci]

* Allow separate inputs to subnets for continuous models (#521)

Introduces easy access to the different inputs x, t and conditions, to allow for specialized processing of each input, which can be beneficial for more advanced use cases.

---------

Co-authored-by: Valentin Pratz <[email protected]>

* Auto-select backend (#543)

* add automatic backend detection and selection

* Fix typo

Co-authored-by: Copilot <[email protected]>

* Add priority ordering of backends

---------

Co-authored-by: Copilot <[email protected]>
Co-authored-by: stefanradev93 <[email protected]>

* Breaking: parameterize MVNormalScore by inverse cholesky factor to improve stability (#545)

* breaking: parameterize MVNormalScore by inverse cholesky factor

The log_prob can be completely calculated using the inverse cholesky
factor L^{-1}. Using this also stabilizes the initial loss, and speeds
up computation.

This commit also contains two optimizations.
Moving the computation of the precision matrix into the einsum, and
using the sum of the logs instead of the log of a product.

As the parameterization changes, this is a breaking change.

* Add right_side_scale_inverse and test [no ci]

The transformation necessary to undo standardization for a Cholesky
factor of the precision matrix is x_ij = x_ij' / sigma_j, which is now
implemented by a right_side_scale_inverse transformation_type.

* Stop skipping MVN tests

* Remove stray keyword argument in fill_triangular_matrix

* Rename cov_chol_inv to precision_chol and update docstrings [no ci]

* rename precision_chol to precision_cholesky_factor

to improve clarity.

* rename cov_chol to covariance_cholesky_factor

* remove check_approximator_multivariate_normal_score function [no ci]

---------

Co-authored-by: han-ol <[email protected]>

* fix unconditional sampling in ContinuousApproximator (#548)

- batch shape was calculated from inference_conditions even if they are
  known to be None
- add approximator test for unconditional setting

* Test quantities Linear Regression Starter notebook (#544)

* Implementation of log-lik test quantity for SBC in starter notebook

* update data-dependent test-quantities example

* Small typo fixes in linear regression notebook

---------

Co-authored-by: Paul-Christian Bürkner <[email protected]>

* fix: optimizer was not used in workflow with multiple fits

For the optimizer to be used, the approximator.compile function has to
be called. This was not the case. I adapted the `setup_optimizer`
function to match the description in its docstring, and made the
compilation conditional on its output. The output indicates if a new
optimizer was configured.

* fix: remove extra deserialize call for SummaryNetwork

The extra call leads to the DTypePolicy to be deserialized. This is then
passed as a class, and cannot be handled by autoconf, leading to the
error discussed in
#549

* Compatibility: deserialize when get_config was overridden

* unify log_prob signature in PointApproximator [no ci]

ContinuousApproximator and BasicWorkflow allow passing the data
positionally, we can allow the same for the PointApproximator.

* Tutorial on spatial data with Gaussian Random Fields (#540) [no ci]

The tutorial uses the experimental ResNet class to build a summary
network for spatial data.

* Support non-array data in test_quantity calibration ecdf [no ci]

Simulator outputs are allowed to be of type int or float, and
consequently have no batch dimension. This needs to be considered
in the broadcasting of inference_conditions for data based SBC
test quantities.

"examples/Linear_Regression_Starter.ipynb" contains an example where this is
necessary, where N is a non-batchable integer.

* import calibration_log_gamma in diagnostics namespace [no ci]

* Add wrapper around scipy.integrate.solve_ivp for integration

* minor fixes and improvements to the pairs plot functions

- pass target color to legend
- do not use common norm, so that prior stays visible in kdeplots
- do not share y on the diagonal, so that all marginal distributions
  stay visible, even if one is very peaked

* fix: layers were not deserialized for Sequential and Residual

As layers were passed with the `*layers` syntax, they could not be
passed as keyword arguments. In `from_config`, however, this was
attempted, leading to the layers to be ignored during reserialization.
This commit fixes this by taking the layers from `kwargs` if they are
passed as a keyword argument.

* add serialization tests for Sequential and Residual

* Fix: ensure checkpoint filepath exists before training

Previously choosing a non-existant directory as checkpoint_filepath
would lead to silently not saving at all.

* Revert 954c16c since it was unnecessary

The alledged issue didn't exist and checkpoint folders are created by
the keras callback automatically already.

I misread tests on this and didn't catch that the problem I was seeing
was caused by a different part of my pipeline.

* improvements to diagnostic plots (#556)

* improvements to diagnostics plots

add markersize parameter, add tests, support dataset_id for
pairs_samples

Fixes #554.

* simplify test_calibration_ecdf_from_quantiles

* Add pairs plot for arbitrary quantities (#550)

Add pairs_quantity and plot_quantity functions that allow plotting of quantities that can be calculated for each individual dataset. Currently, for the provided metrics this is only useful for posterior contraction, but could be useful for posterior z-score and other quantities as well.

* minor fix in diffusion edm schedule (#560)

* minor fix in diffusion edm schedule

* DeepSet: Adapt output dimension of invariant module inside the equivariant module (#557) (#561)

* adapt output dim of invariant module in equivariant module

See #557. The DeepSet showed bad performance and was not able to learn
diverse summary statistics. Reducing the dimension of the output of the
invariant module inside the equivariant module improves this, probably
because the invidividual information of each set member gains importance
compared to the shared information provided by the invariant module.

There might be better settings for this, so we might update the default
later on. However, this is already an improvement over the previous
setting.

* DeepSet: adapt docstring to reflect code

* pairs_postorior: inconsistent type hint fix (#562)

* allow exploding variance type in EDM schedule

* fix type hint

* Bump up version [skip ci]

* Fix instructions for backend spec [skip ci]

* Add New Flow Matching Schedules (#565)

* add fm schedule

* add fm schedule

* add comments

* expose time_power_law_alpha

* Improve doc [skip ci]

---------

Co-authored-by: stefanradev93 <[email protected]>

* change default integration method to rk45

for DiffusionModel and FlowMatching. Euler shows significant deviations
when computing the log-prob, which risks misleading users regarding the
performance of the networks.

rk45 is slower, but the problem is heavily reduced with this method.

* fix nan to num inverse

* fix setting markersize in lotka volterra notebook

* fix: actually set KERAS_BACKEND to chosen backend

Add warning if KERAS_BACKEND and actually loaded backend do not match.
This can happen if keras is imported before BayesFlow.

* Fix warning msg

---------

Co-authored-by: Valentin Pratz <[email protected]>
Co-authored-by: han-ol <[email protected]>
Co-authored-by: Daniel Habermann <[email protected]>
Co-authored-by: Valentin Pratz <[email protected]>
Co-authored-by: arrjon <[email protected]>
Co-authored-by: Lars <[email protected]>
Co-authored-by: Hans Olischläger <[email protected]>
Co-authored-by: Lasse Elsemüller <[email protected]>
Co-authored-by: Leona Odole <[email protected]>
Co-authored-by: Jonas Arruda <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Paul-Christian Bürkner <[email protected]>
Co-authored-by: The-Gia Leo Nguyen <[email protected]>

* fix trainable parameters in distributions (#520)

* Improve numerical precision in MVNScore.log_prob

* add log_gamma diagnostic (#522)

* add log_gamma diagnostic

* add missing export for log_gamma

* add missing export for gamma_null_distribution, gamma_discrepancy

* fix broken unit tests

* rename log_gamma module to sbc

* add test_log_gamma unit test

* add return information to log_gamma doc string

* fix typo in docstring, use fixed-length np array to collect log_gammas instead of appending to an empty list

* Breaking changes: Fix bugs regarding counts in standardization layer (#525)

* standardization: add test for multi-input values (failing)

This test reveals to bugs in the standarization layer:

- count is updated multiple times
- batch_count is too small, as the sizes from reduce_axes have to be
  multiplied

* breaking: fix bugs regarding count in standardization layer

Fixes #524

This fixes the two bugs described in c4cc133:

- count was accidentally updated, leading to wrong values
- count was calculated wrongly, as only the batch size was used. Correct
  is the product of all reduce dimensions. This lead to wrong standard
  deviations

While the batch dimension is the same for all inputs, the size of the
second dimension might vary. For this reason, we need to introduce an
input-specific `count` variable. This breaks serialization.

* fix assert statement in test

* rename log_gamma to calibration_log_gamma (#527)

* simple fix

* Hotfix: numercial stability of non-log-stabilized sinkhorn plan (#531)

* fix numerical stability issues in sinkhorn plan

* improve test suite

* fix ultra-strict convergence criterion in log_sinkhorn_plan

* update dependencies

* add comment about convergence check

* update docsting to reflect fixes

* sinkhorn_plan now returns a transport plan with uniform marginal distributions

* add unit test for sinkhorn_plan

* fix sinkhorn function by sampling from the logits of the transpose of the plan, instead of the plan directly

* sinkhorn(x1, x2) now samples from log(plan) to receive assignments such that x2[assignments] matches x1

* re-enable test_assignment_is_optimal() for method='sinkhorn'

* log_sinkhorn now correctly uses log_plan instead of keras.ops.exp(log_plan), log_sinkhorn_plan returns logits of the transport plan

* add unit tests for log_sinkhorn_plan

* fix faulty indexing with tensor for tensorflow backend

* re-add numItermax for ot pot test

---------

Co-authored-by: Daniel Habermann <[email protected]>

* isinstance sequence

* Pass correct training stage in compute_metrics (#534)

* Pass correct training stage in CouplingFlow.compute_metrics

* Pass correct training stage in CIF and PointInferenceNetwork

* Custom test quantity support for calibration_ecdf (#528)

* Custom test quantity support for calibration_ecdf

* rename variable [no ci]

* Consistent defaults for variable_keys/names in calibration_ecdf with test quantiles

* Tests for calibration_ecdf with test_quantities

* Remove redundant and simplify comments

* Fix docstrings and typehints

---------

Co-authored-by: stefanradev93 <[email protected]>

* Log gamma test fix (#535)

* fix test_calibration_log_gamma_end_to_end unit test failing too often than expected

* set alpha to 0.1% in binom.ppf

* fix typo in comment

* Stateless adapters (#536)

* Remove stateful adapter features

* Fix tests

* Fix typo

* Remove nnpe from adapter

* Bring back notes [skip ci]

* Remove unncessary restriction to kwargs only [skip ci]

* Remove old super call [skip ci]

* Robustify type [skip ci]

* remove standardize from multimodal sim notebook [no ci]

* add draft module docstring to augmentations module [no ci]

Feel free to modify.

* adapt and run neurocognitive modeling notebook [no ci]

* adapt cCM playground notebook [no ci]

* adapt signature of Adapter.standardize

* add parameters missed in previous commit

* Minor NNPE polishing

* remove stage in docstring from OnlineDataset

---------

Co-authored-by: Lasse Elsemüller <[email protected]>
Co-authored-by: Valentin Pratz <[email protected]>

* Fix training strategies in BasicWorkflow

* move multimodal data notebook to regular examples [no ci]

* make pip install call on homepage more verbose [no ci]

* remove deprecated summaries function

The function was renamed to summarize in v2.0.4.

* detail subsampling behavior docs for SIR simulator [no ci]

fixes #518

* move DiffusionModel from experimental to networks

Stabilizes the DiffusionModel class. A deprecation warning for the
DiffusionModel class in the experimental module was added.

* Add citation for resnet (#537) [no ci]

* added citation for resnet

* minor formatting

---------

Co-authored-by: Valentin Pratz <[email protected]>

* Bump up version [skip ci]

* Allow separate inputs to subnets for continuous models (#521)

Introduces easy access to the different inputs x, t and conditions, to allow for specialized processing of each input, which can be beneficial for more advanced use cases.

---------

Co-authored-by: Valentin Pratz <[email protected]>

* Auto-select backend (#543)

* add automatic backend detection and selection

* Fix typo

Co-authored-by: Copilot <[email protected]>

* Add priority ordering of backends

---------

Co-authored-by: Copilot <[email protected]>
Co-authored-by: stefanradev93 <[email protected]>

* Breaking: parameterize MVNormalScore by inverse cholesky factor to improve stability (#545)

* breaking: parameterize MVNormalScore by inverse cholesky factor

The log_prob can be completely calculated using the inverse cholesky
factor L^{-1}. Using this also stabilizes the initial loss, and speeds
up computation.

This commit also contains two optimizations.
Moving the computation of the precision matrix into the einsum, and
using the sum of the logs instead of the log of a product.

As the parameterization changes, this is a breaking change.

* Add right_side_scale_inverse and test [no ci]

The transformation necessary to undo standardization for a Cholesky
factor of the precision matrix is x_ij = x_ij' / sigma_j, which is now
implemented by a right_side_scale_inverse transformation_type.

* Stop skipping MVN tests

* Remove stray keyword argument in fill_triangular_matrix

* Rename cov_chol_inv to precision_chol and update docstrings [no ci]

* rename precision_chol to precision_cholesky_factor

to improve clarity.

* rename cov_chol to covariance_cholesky_factor

* remove check_approximator_multivariate_normal_score function [no ci]

---------

Co-authored-by: han-ol <[email protected]>

* fix unconditional sampling in ContinuousApproximator (#548)

- batch shape was calculated from inference_conditions even if they are
  known to be None
- add approximator test for unconditional setting

* Test quantities Linear Regression Starter notebook (#544)

* Implementation of log-lik test quantity for SBC in starter notebook

* update data-dependent test-quantities example

* Small typo fixes in linear regression notebook

---------

Co-authored-by: Paul-Christian Bürkner <[email protected]>

* fix: optimizer was not used in workflow with multiple fits

For the optimizer to be used, the approximator.compile function has to
be called. This was not the case. I adapted the `setup_optimizer`
function to match the description in its docstring, and made the
compilation conditional on its output. The output indicates if a new
optimizer was configured.

* fix: remove extra deserialize call for SummaryNetwork

The extra call leads to the DTypePolicy to be deserialized. This is then
passed as a class, and cannot be handled by autoconf, leading to the
error discussed in
#549

* Compatibility: deserialize when get_config was overridden

* unify log_prob signature in PointApproximator [no ci]

ContinuousApproximator and BasicWorkflow allow passing the data
positionally, we can allow the same for the PointApproximator.

* Tutorial on spatial data with Gaussian Random Fields (#540) [no ci]

The tutorial uses the experimental ResNet class to build a summary
network for spatial data.

* Support non-array data in test_quantity calibration ecdf [no ci]

Simulator outputs are allowed to be of type int or float, and
consequently have no batch dimension. This needs to be considered
in the broadcasting of inference_conditions for data based SBC
test quantities.

"examples/Linear_Regression_Starter.ipynb" contains an example where this is
necessary, where N is a non-batchable integer.

* import calibration_log_gamma in diagnostics namespace [no ci]

* Add wrapper around scipy.integrate.solve_ivp for integration

* minor fixes and improvements to the pairs plot functions

- pass target color to legend
- do not use common norm, so that prior stays visible in kdeplots
- do not share y on the diagonal, so that all marginal distributions
  stay visible, even if one is very peaked

* fix: layers were not deserialized for Sequential and Residual

As layers were passed with the `*layers` syntax, they could not be
passed as keyword arguments. In `from_config`, however, this was
attempted, leading to the layers to be ignored during reserialization.
This commit fixes this by taking the layers from `kwargs` if they are
passed as a keyword argument.

* add serialization tests for Sequential and Residual

* Fix: ensure checkpoint filepath exists before training

Previously choosing a non-existant directory as checkpoint_filepath
would lead to silently not saving at all.

* Revert 954c16c since it was unnecessary

The alledged issue didn't exist and checkpoint folders are created by
the keras callback automatically already.

I misread tests on this and didn't catch that the problem I was seeing
was caused by a different part of my pipeline.

* improvements to diagnostic plots (#556)

* improvements to diagnostics plots

add markersize parameter, add tests, support dataset_id for
pairs_samples

Fixes #554.

* simplify test_calibration_ecdf_from_quantiles

* Add pairs plot for arbitrary quantities (#550)

Add pairs_quantity and plot_quantity functions that allow plotting of quantities that can be calculated for each individual dataset. Currently, for the provided metrics this is only useful for posterior contraction, but could be useful for posterior z-score and other quantities as well.

* minor fix in diffusion edm schedule (#560)

* minor fix in diffusion edm schedule

* DeepSet: Adapt output dimension of invariant module inside the equivariant module (#557) (#561)

* adapt output dim of invariant module in equivariant module

See #557. The DeepSet showed bad performance and was not able to learn
diverse summary statistics. Reducing the dimension of the output of the
invariant module inside the equivariant module improves this, probably
because the invidividual information of each set member gains importance
compared to the shared information provided by the invariant module.

There might be better settings for this, so we might update the default
later on. However, this is already an improvement over the previous
setting.

* DeepSet: adapt docstring to reflect code

* pairs_postorior: inconsistent type hint fix (#562)

* allow exploding variance type in EDM schedule

* fix type hint

* Bump up version [skip ci]

* Fix instructions for backend spec [skip ci]

* Add New Flow Matching Schedules (#565)

* add fm schedule

* add fm schedule

* add comments

* expose time_power_law_alpha

* Improve doc [skip ci]

---------

Co-authored-by: stefanradev93 <[email protected]>

* change default integration method to rk45

for DiffusionModel and FlowMatching. Euler shows significant deviations
when computing the log-prob, which risks misleading users regarding the
performance of the networks.

rk45 is slower, but the problem is heavily reduced with this method.

* fix nan to num inverse

* fix setting markersize in lotka volterra notebook

* fix: actually set KERAS_BACKEND to chosen backend

Add warning if KERAS_BACKEND and actually loaded backend do not match.
This can happen if keras is imported before BayesFlow.

* Fix warning msg

* Adapt sCMs [skip ci]

* Update exp notebooks [skip ci]

* Enable Python 3.13 support, allow numpy 2.x (#568)

* Enable Python 3.13 support, allow numpy 2.x

* Disable Python 3.13 test again [no ci]

* Remove deprecated CHANGELOG.rst [skip ci] (#570)

* add parallel simulator (#569)

* allow tensor in DiagonalNormal dimension (#571)

* allow tensor in DiagonalNormal dimension

* Update student and mixture with multiple dims too

* Use credible intervals for uncertainties in recovery plots (#573)

* use credible intervals for uncertainties in recovery plots

* Add kwargs for agg functions

* make point_arg_kwargs and uncertainty_agg_kwargs explicit arguments

* adapt docs, minor stylistic changes

* add support for symmetric uncertainty measures

* add test for symmetric uncertainty measure in recovery

---------

Co-authored-by: stefanradev93 <[email protected]>
Co-authored-by: Valentin Pratz <[email protected]>

* fix serialization in StableConsistencyModel (#578)

* Fix SetTransformer kwargs handling (#574)

* fix kwargs in sample and set transformer

* fix scm

* Add empirical coverage diagnostic plots (#579)

* add empirical coverage diagnostic plots

* merging the plotting functions via difference arg and add beta binomial reference

* adjust tests and diagnostics init

* Allow different normalizations for NRMSE

* add missing standardization calls to Approximator.summarize (#581)

* relax test_calibration_log_gamma_end_to_end() by introducing a larger posterior bias and decreasing the alpha error rate to 0.01% (#582)

* Improve default diagnostics (#584)

* Add coverage to default diagnostics

* Improve diagnostics

* Adapt test

* Account for new metric

* Fix doc [skip ci]

* Adding custom test quantities to more diagnostics (#586)

* Adding custom test quantities to diagnostics

* Fix Code Style

* run ruff again

* Adding tests

* adding one stupid blank line

* again run ruff

---------

Co-authored-by: Paul-Christian Bürkner <[email protected]>

* fix scm (#590)

* Fix stable consistency (#595)

* fix scm

* fix saving

* Adding z-score metric and fixing test_quantities problem in plot_quantity/pairs_quantity (#596)

* Add posterior z score

* Fixing test_quantities problem in plot_quantity/pairs_quantity

* Removing unused variable found by code style check

* Fix benchmark (#597)

* fix sir

* fix lotka volterra

* fix lotka volterra

* fix sir test

* fix LV

* Add cross-validation support to classifier_two_sample_test (#600)

* add optional time embedding layer to consistency model (#601)

* fix: add optional time embedding layer to consistency model

* Simplify [skip ci]

---------

Co-authored-by: stefanradev93 <[email protected]>

* Refactoring of integration routines (#598)

* correct rk45 and add tsit5

* add predictor corrector

* add adaptive sampler SDE

* add shark

* rm warn

* fix dt

* fix adaptive step size

* refactor stochastic integrator

* refactor stochastic integrator

* refactor stochastic integrator

* fix adaptive

* fix Tsit5

* fix sampler

* updated stochastic solvers

* add Langevin

* add Langevin

* add adaptive step size

* tune adaptive step size

* add Gotta Go Fast SDE sampler

* improve adaptive ODE samplers

* fix schedule test

* improved defaults

* improved defaults

* improved initial step size

* improved initial step size

* check nan in integrate

* set default

* update model defaults

* make loop jax compatible

* filter kwargs

* fix density computation

* fix jax all nans

* fix jax all nans

* fix jax all nans

* relax tols in tests

* enable density computation with adaptive step size solvers

* update max_steps parameter for integration to reduce memory usage

* update noise handling in integration methods

* add jax support for adaptive sampling

* more informative warnings

* more informative warnings

* count_not_accepted

* improved tsit5 on gpu

* improved tsit5 on gpu

* improved tsit5 on gpu

* improved tsit5 on gpu

* Small consistency tweaks

---------

Co-authored-by: stefanradev93 <[email protected]>

* automatic MLP widths in c2st (#607)

* automatic MLP widths in c2st

* update doc

* Add conditional optimal transport for flow matching (#599)

* Add partial optimal transport support to sinkhorn and log_sinkhorn functions

* improve type hints

* optimal transport with conditional support

* improve regularization handling

* fix regularization

* fix if

* refactor optimal transport to include conditions in return values

* improve warnings

* improve warnings

* improve warnings

* improve tests

* improve tests

* improved keywords

* improved partial costs

* faster implementation

* Fix naming convention and docs

* remove assignments

* update docs

* rename conditional_ot_ratio to condition_ratio in default dict

---------

Co-authored-by: stefanradev93 <[email protected]>

* fix output projector initialization in single coupling (#610)

* fix output projector initialization in single coupling

* Remove units=None and move projector creation to build

* Fix output projector initialization to use 'units' parameter consistently

* Remove kernel_initializer from output projector layers

---------

Co-authored-by: stefanradev93 <[email protected]>

* Inference without an additional sample dimension (#611)

* Remove additional sample dim internally

* Do not reshape for unconditional estimation

* Fix point approximator issues

* Simplify

* Time mlp (#612)

* add TimeMLP and time embedding support to MLP

* fix time embedding dimension calculation in stable consistency model

* add ConditionalResidual and FiLM layers to TimeMLP; refactor time embedding handling

* fix build

* fix build in FiLM

* fix build

* fix build

* fix build

* add super build

* fix fourier embedding

* fix fourier embedding

* fix fourier embedding

* better defaults

* better defaults

* remove film

* remove film

* better defaults

* better defaults

* better defaults

* add film

* fix dims

* fix kwargs

* separate gamma and beta projections

* fix test

* fix test

* change default subnet in diffusion type models

* remove unnecessary dimension expansion for time input

* Add variance_type to config in edm_noise_schedule

* fix merge conflict with subnet projector

* Enhance documentation

* change default concatenation

* update docstring

* fix test

* fix defaults in test

* Prepare for v2.0.8 [skip ci]

* Fix docs

* Cleanup experimental module and pull out stable consistency

* Cleanup experimental module and pull out stable consistency

* Fix import

* Fix import

* Cleanup subnet handling in inference networks (#614)

* Cleanup subnet handling in inference networks

* fix test

* Cleanup and bugfix partial_ot_factor -> partial_factor

---------

Co-authored-by: stefanradev93 <[email protected]>

* Add batch sampling to workflow (#609)

* add training duration logging to basic workflow

* add batch sampling support to workflow

* add progress bar for batch sampling

* Simplify sampling logic and code, fix progress bar, use existing utils, remove vstack_samples

* Ensure everything works with nested dicts / fusion networks

* Add diffusion / consistency to test workflows, reduce num epochs

* fix test

---------

Co-authored-by: stefanradev93 <[email protected]>

* Flexible sample shapes for time series, images, sets, etc. (#616)

* Add new sample_shape argument

* Relax default with no conditions

* Fix @*!@%$!* inconsistency in backend behavior

* Fix @*!@%$!* inconsistency in backend behavior

* Fix dims for diagonal student

* Fix density eval adaptive solvers (#617)

* Update ODE solver with adaptive steps for density calculation

* Update ODE solver with adaptive steps for density calculation

* Update ODE solver with adaptive steps for density calculation

* more time logging

* fix test for jax

* fix test for jax

* fix notebooks

* always broadcast time in flow matching

* remove training kwarg from output projector

* add warnings for stochastic methods in forward integration and density computation

* Slight improvement to error message [skip ci]

---------

Co-authored-by: stefanradev93 <[email protected]>

* Enable diffusion guidance (#618)

* Add support for custom guidance in diffusion model

* Add support for torch

* Fix naming and remove redundant training arg to projector

* Fix num samples across the test

---------

Co-authored-by: stefanradev93 <[email protected]>

* Update notebooks

* Fix warning message [skip ci]

* Always use default integration method for DiffusionModels in tests

* Fix list to string concatenation

* add refs [skip ci]

* upgrade dependencies and refactor optional [all] dependency (#620)

* add error msg in test

* Fix SIR notebook

* Update landing fig (#621)

* Update landing fig

* Streamline desc [skip ci]

---------

Co-authored-by: stefanradev93 <[email protected]>

* fix seed in density computation

* Bring down rtol for numerical J

* improve test density computation

* remove unnecessary kwargs

* Update test workflow to only trigger slow tests on JAX backend
also greatly improve readability of the if statement

---------

Co-authored-by: Valentin Pratz <[email protected]>
Co-authored-by: han-ol <[email protected]>
Co-authored-by: Daniel Habermann <[email protected]>
Co-authored-by: Valentin Pratz <[email protected]>
Co-authored-by: arrjon <[email protected]>
Co-authored-by: Lars <[email protected]>
Co-authored-by: Hans Olischläger <[email protected]>
Co-authored-by: Lasse Elsemüller <[email protected]>
Co-authored-by: Leona Odole <[email protected]>
Co-authored-by: Jonas Arruda <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Paul-Christian Bürkner <[email protected]>
Co-authored-by: The-Gia Leo Nguyen <[email protected]>
Co-authored-by: Lasse Meixner <[email protected]>
Co-authored-by: Svenja Jedhoff <[email protected]>
Co-authored-by: Jerry Huang <[email protected]>