[Feature] Add Latent Diffusion and generic latent inference wrapper for all inference networks. by bhanuprasanna2001 · Pull Request #656 · bayesflow-org/bayesflow

bhanuprasanna2001 · 2026-03-17T21:37:52Z

Included the use of any inference network as part of the latent diffusion model.
Only remaining task is to update the model's name to a more suitable one.
Developed a notebook experiment with two moons to test all inference networks and observe their behavior in the latent space.
The experiment is functioning well.

…eusable

…sionModel

Copilot

Pull request overview

Adds a new latent-diffusion-based inference network to BayesFlow by combining a VAE-style encoder/decoder with an existing latent-space inference network (DiffusionModel by default), along with a dedicated pytest suite and public exports.

Changes:

Introduces LatentDiffusionModel plus standalone Encoder/Decoder components under bayesflow.networks.inference.latent_diffusion.
Exposes the new classes via bayesflow.networks.inference and bayesflow.networks package exports.
Adds comprehensive tests for build/shape/metrics/serialization and flow-matching integration.

Reviewed changes

Copilot reviewed 10 out of 11 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
`bayesflow/networks/inference/latent_diffusion/latent_diffusion_model.py`	Implements the composed latent diffusion inference model and training objective.
`bayesflow/networks/inference/latent_diffusion/encoder.py`	Adds a VAE-style encoder producing `(z, mean, log_var)` with reparameterization.
`bayesflow/networks/inference/latent_diffusion/decoder.py`	Adds a decoder mapping latent vectors back to the original parameter space.
`bayesflow/networks/inference/latent_diffusion/__init__.py`	Registers and exports latent diffusion components for the subpackage.
`bayesflow/networks/inference/__init__.py`	Re-exports `LatentDiffusionModel`, `Encoder`, and `Decoder` at the inference package level.
`bayesflow/networks/__init__.py`	Re-exports the new inference components at the top-level networks package.
`tests/test_networks/test_latent_diffusion_model/conftest.py`	Adds fixtures for latent diffusion model tests (dims, conditions, model variants).
`tests/test_networks/test_latent_diffusion_model/test_latent_diffusion_model.py`	Tests LDM build, forward/inverse shapes, metrics, warmup, serialization, save/load.
`tests/test_networks/test_latent_diffusion_model/test_encoder.py`	Tests encoder build/output shapes, auto latent dim, stochasticity, serialization, save/load.
`tests/test_networks/test_latent_diffusion_model/test_decoder.py`	Tests decoder build/output shapes, required `output_dim`, serialization, save/load.
`tests/test_networks/test_latent_diffusion_model/__init__.py`	Declares the new test package.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

+        # Build encoder
+        self.encoder.build(xz_shape)
+        actual_latent_dim = self.encoder.latent_dim
+
+        # Build decoder with correct output dimension
+        self.decoder.output_dim = input_dim
+        latent_shape = tuple(xz_shape[:-1]) + (actual_latent_dim,)
+        self.decoder.build(latent_shape)


+        Network for diffusion noise prediction (only used when
+        ``inference_network`` is None). Default is ``"time_mlp"``.
+    diffusion_subnet_kwargs : dict[str, any], optional
+        Additional arguments for diffusion subnet (only used when
+        ``inference_network`` is None). Default is None.
+    noise_schedule : str, optional
+        Noise schedule for diffusion (only used when ``inference_network``
+        is None). Default is ``"cosine"``.
+    schedule_kwargs : dict[str, any], optional
+        Additional arguments for noise schedule (only used when
+        ``inference_network`` is None). Default is None.
+    integrate_kwargs : dict[str, any], optional
+        Configuration for ODE integration during sampling (only used when
+        ``inference_network`` is None). Default is None.


+        elif isinstance(inference_network, str):
+            raise ValueError(
+            f"Unknown inference_network specification: {inference_network}. "
+            f"Expected 'auto' or an InferenceNetwork instance."
+            )
+        else:
+            self.inference_network = inference_network


+def test_build(latent_diffusion_model, random_samples, random_conditions):
+    xz_shape = keras.ops.shape(random_samples)
+    cond_shape = keras.ops.shape(random_conditions) if random_conditions is not None else None
+
+    assert not latent_diffusion_model.built
+    latent_diffusion_model.build(xz_shape, conditions_shape=cond_shape)
+    assert latent_diffusion_model.built
+    assert latent_diffusion_model.variables
+
+
+def test_forward_output_shape(latent_diffusion_model, random_samples, random_conditions):
+    xz_shape = keras.ops.shape(random_samples)
+    cond_shape = keras.ops.shape(random_conditions) if random_conditions is not None else None
+    latent_diffusion_model.build(xz_shape, conditions_shape=cond_shape)
+
+    z = latent_diffusion_model(random_samples, conditions=random_conditions)
+    assert keras.ops.shape(z) == (xz_shape[0], latent_diffusion_model.latent_dim)
+


@@ -0,0 +1,84 @@
+import keras
+import numpy as np
+import pytest


@@ -0,0 +1,66 @@
+import keras
+import numpy as np


bhanuprasanna2001

Added Latent Inference Network, which can be used with any Inference Network to operate in the latent space.
Added Latent Diffusion Model as a thin wrapper for the Latent Inference Network.

…orflow backend

…(jvp, vjp) to autodiff module

… jvp, vjp, and jacobians

…inary operations

…utograd

Co-authored-by: Copilot <copilot@github.com>

bhanuprasanna2001 added 2 commits March 17, 2026 19:01

refactor: LDM - separate encoder/decoder, fix bugs, add tests, make r…

bc4bc41

…eusable

refactor: update imports and improve docstring clarity in LatentDiffu…

7712cfb

…sionModel

bhanuprasanna2001 marked this pull request as ready for review March 17, 2026 21:37

Copilot AI review requested due to automatic review settings March 17, 2026 21:38

Copilot started reviewing on behalf of bhanuprasanna2001 March 17, 2026 21:38 View session

Copilot AI reviewed Mar 17, 2026

View reviewed changes

LarsKue assigned bhanuprasanna2001 Mar 18, 2026

LarsKue added the feature New feature or request label Mar 18, 2026

LarsKue added this to the BayesFlow 2.1 milestone Mar 18, 2026

LarsKue changed the title ~~Latent Space support for all inference networks. Tested.~~ [Feature] Add Latent Diffusion and generic latent inference wrapper for all inference networks. Mar 18, 2026

Refactor latent inference network components and tests

f019e52

bhanuprasanna2001 commented Apr 1, 2026

View reviewed changes

LarsKue added 18 commits April 16, 2026 17:27

add backend autograd functional api

64ed632

Fix gradient extraction for single argument case in autograd for tens…

b4f912b

…orflow backend

add backend-agnostic jit

d5040ab

add intermediate version of trainer

ef746aa

add training example

723e6e5

add jacobian functions (jacfwd, jacrev) and vector-Jacobian products …

f7c3150

…(jvp, vjp) to autodiff module

add unit tests for autodiff functions including grad, value_and_grad,…

760fc10

… jvp, vjp, and jacobians

add extensive unit tests for autodiff functions including unary and b…

16af2c9

…inary operations

fix torch backend jvp to ensure primals and tangents are always tuples

5ad9b21

[wip] fix tensorflow autodiff implementations

03ff91a

make return type tests for jacfwd and jacrev more concrete

d1c7e45

remove commented code

dce5b61

update backend autodiff tests to include jit compilation

fe9f6be

fix shape for scalar jacobian in tensorflow

69a2915

change torch jit backend to inductor again

9b7fa4b

fix error message formatting for backend setup failure

d27e28b

Make backend module protected

b6c3d6d

Remove intermediate-state trainers module

d22f00a

LarsKue and others added 10 commits April 24, 2026 15:42

Remove intermediate-state trainers module and example

b61923f

Merge remote-tracking branch 'origin/backend-autograd' into backend-a…

40c038e

…utograd

add jit compilation to base jacfwd and jacrev tests

f7b21e7

fix value_and_grad and grad for tensorflow backend with jit compilation

29979ca

fix: use bound_fn in value_and_grad

8ad5f5c

split jacfwd and jacrev tests into subtests for better error diagnostics

73a2b39

add vibe-coded version of jacfwd and jacrev that passes all tests

b91a935

Merge upstream/backend-autograd into ldm-dev

8c02e32

Remove ldm

d51917f

deterministic ae.

07e624f

Co-authored-by: Copilot <copilot@github.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Feature] Add Latent Diffusion and generic latent inference wrapper for all inference networks.#656

[Feature] Add Latent Diffusion and generic latent inference wrapper for all inference networks.#656
bhanuprasanna2001 wants to merge 31 commits into
bayesflow-org:devfrom
bhanuprasanna2001:ldm-dev

bhanuprasanna2001 commented Mar 17, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

bhanuprasanna2001 left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

bhanuprasanna2001 commented Mar 17, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

bhanuprasanna2001 left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants