restructuring neural models + addition of OT-FM and GENOT #468

MUCDK · 2023-11-22T09:05:04Z

This is a PR for

new base classes neural solvers and models (i.e. neural networks)
Incorporating unbalancedness and learning the rescaling factors for any neural OT model.
adding OTFlowMatching and, related to this, classes for flows and time samplers
adding GENOT (with extension to conditional GENOT)
adding drafts of data loaders.

Following this PR, the implementations of ICNN-based solvers and the Monge Gap model should be adapted and extended to the unbalanced setting.

~~Moreover, wrt typing, I replaced jnp.ndarray by jax.Array~~

What remains to be done, but I would prefer to do in a separate PR

add graph costs to OTFM and GENOT, i.e. functions which compute batch-wise graphs, and compute costs, e.g. geodesic Sinkhorn or convolutional Wasserstein from this.
implementations of ICNN-based solvers and the Monge Gap model should be adapted and extended to the unbalanced setting.

codecov · 2023-11-22T10:37:52Z

Codecov Report

Attention: Patch coverage is 91.97652% with 41 lines in your changes are missing coverage. Please review.

Project coverage is 90.71%. Comparing base (41906a2) to head (f227d54).
Report is 1 commits behind head on main.

❗ Current head f227d54 differs from pull request most recent head 6f9a77c. Consider uploading reports for the commit 6f9a77c to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #468      +/-   ##
==========================================
+ Coverage   90.64%   90.71%   +0.07%     
==========================================
  Files          60       68       +8     
  Lines        6682     7046     +364     
  Branches      956      996      +40     
==========================================
+ Hits         6057     6392     +335     
- Misses        477      494      +17     
- Partials      148      160      +12

Files	Coverage Δ
src/ott/datasets.py	`95.00% <100.00%> (ø)`
src/ott/initializers/neural/meta_initializer.py	`92.42% <100.00%> (ø)`
src/ott/neural/methods/monge_gap.py	`91.20% <100.00%> (ø)`
src/ott/neural/methods/neuraldual.py	`57.14% <100.00%> (ø)`
src/ott/neural/networks/layers/conjugate.py	`100.00% <ø> (ø)`
src/ott/neural/networks/layers/posdef.py	`91.30% <ø> (ø)`
src/ott/neural/networks/layers/time_encoder.py	`100.00% <100.00%> (ø)`
src/ott/neural/networks/velocity_field.py	`100.00% <100.00%> (ø)`
src/ott/problems/linear/potentials.py	`91.44% <100.00%> (ø)`
src/ott/solvers/linear/lineax_implicit.py	`100.00% <100.00%> (ø)`
... and 8 more

... and 1 file with indirect coverage changes

MUCDK · 2023-11-22T13:36:06Z

co-authored by @lucaeyring

…Array

review-notebook-app · 2023-11-26T14:57:42Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

…olver

JTT94 · 2024-03-25T17:37:25Z

src/ott/neural/flow_models/flows.py

+ return jnp.full_like(t, fill_value=self.sigma)
+
+
+class BrownianNoiseFlow(StraightFlow):


"BrownianNoiseFlow" is ambiguous to me, it could correspond to Brownian motion and associated flow e.g. VE Flow in eqn 16 of https://arxiv.org/pdf/2210.02747.pdf

src/ott/neural/flow_models/genot.py

src/ott/neural/flow_models/flows.py

michalk8 · 2024-03-29T18:19:06Z

TODOs left:

check the docstring for wrong links
check imports in all the neural tutorials
update the documentation

MUCDK · 2024-04-02T08:36:16Z

tests/neural/conftest.py

+
+
+@pytest.fixture()
+def lin_cond_dl() -> DataLoader:


This is not a "conditional OT data loader".

src/ott/neural/methods/flows/genot.py

michalk8

Thanks @MUCDK and everyone, LGTM, merging this!

* draft of BaseSolver and UnbalancedMixin * draft of BaseSolver and UnbalancedMixin * [ci skip] continue flow matching implementation * [ci skip] continue flow matching implementation * [ci skip] add neural networks * [ci skip] add test * [ci skip] resolve import errors * [ci skip] MRO not working * [ci skip] basic test for flow matching passes * [ci skip] add tests for FM with conditions and conditional OT with FM * [ci skip] add genot outline * [ci skip] restructure genot * [ci skip] restructure genot * [ci skip] fix transport * [ci skip] flow matching tests passing * [ci skip] add more tests genot * [ci skip] add more tests genot * [ci skip] add TimeSampler * [ci skip] add docs for TimeSampler and Flow * [ci skip] add docs for OTFlowMatching and replace jnp.ndarray by jax.Array * [ci skip] change init arguments of GENOT and add docstrings to GENOT * [ci skip] split nets into base_models and models * [ci skip] add references * add tests for learning the rescaling factors * [ci skip] partially fix rescaling factor learning * [ci skip] fix rescaling factor learning * [ci skip] all tests passing but k_samples_per_x in genot * k_samples_per_x working in GENOT * [ci skip] changed dataloaders to numpy and dict return * [ci skip] changed dataloaders to numpy and dict return * revert jax.Array to jnp.ndarray * move dataloader from tests to module * add docstrings to neurcal networks * [ci skip] adapt type of scale_cost and cost_fn * [ci skip] clean code * [ci skip] fix genot tests * [ci skip] fix otfm tests * [ci skip] fix otfm tests * add scale cost to otfm * incorporate feedback partially * resolve circular import errors * resolve a few pre-commit errors * resolve pre-commit errors * resolve pre-commit errors * fix rng bug * Update pre-commit * fix import error * Run linter * replace rng jnp.ndarray type by jax.array * replace rng jnp.ndarray type by jax.array * fix import error * [ci skip] start to incorporate feedback * restructure neural module * fix import errors * incorporate feedback partially * make time encoder a layer * make conditions Optional and minor feedback * revert faulty jax.array / jnp.ndarray conversions * make formatting in neural nets nicer * add description to Velocity Field * replace time sampler class by function * add citations * add more references * rename keys_model to rng * fix tests regarding time sampling * fix typo in tests * rename neural_vector_field to velocity_field everywhere * fix OTFlowMatching.transport * fix rescaling networks * Update src/ott/neural/flows/flows.py Co-authored-by: nvesseron <[email protected]> * Update src/ott/neural/flows/flows.py Co-authored-by: nvesseron <[email protected]> * test for scale_cost * update test for scale_cost * fix bug for scale_cost * fix bug for scale_cost * jit solve_ode in genot * incorporate changes partially * [ci skip] intermediate save * [ci skip] neural base solver update * make resamlpemixin a class * incorporate more changes * move noise sampling to flows * fix bug in passing rngs in otfm * introduce otmatcher in otfm * [ci skip] split GENOT into GENOTLin and GENOTQuad * remove dictionaries in OTFM and GENOT classes * change logic in match_latent_to_data in genot * change data loaders / data sets * finish data loader refactoring * Update linter * fix bug in _resample_data` * incorporate more changes * add docs * incorporate more changes * problem with custom type * fix scale cost bug * fix bugs * fux bug in unbalancedness/rescalingMlp * unify unbalancedness step in GENOT * change OTDataSet and OTFlowMatching to 4 data loaderes * Fix bug in the `ConditionalOTDataset` * Polish docs in the `flows.py` * Update `OTFM` * Fix small bugs in `OTFM` * Polish layers * Fix typo in citation * More polish for the docs * remove print statements and unbalancednesshandler * remove tests * make genot training loops more similar to otfm training loop * adapt tests to the extent possible * Add weights to sampling * Start cleaning matchers * Add conditional sampling + resampling * Add initial quad matcher * Improve typing * Remove `base_solver.py` * Add TODO * Update datasets, fix OTFM tests * Start cleaning GENOT * Update GENOT * Remove old GENOTLin/GENOTQuad * Remove axis swapping * Remove old todo * Fix OTFM tests * Remove `MLPBlock` and `RescalingMLP` * Add forgotten license * Remove `__post_init__` from `VF` * Move cyclical time encoder * Move more stuff to `utils` * Remove `samplers.py` * Rename `cond_dim` -> `condition_dim` * Nicer formatting * Fix bug when sampling from the target * Fix another bug when sampling from the data * Add initial test for GW * Remove old GENOT tests * Remove old dataloaders * Add more todos * add docs to dataloader * expose args in GENOT * add docs and adapt data_match_fn * fix linting * fix data loading and add genot fused tests * genot tests passing * adapt docs * adapt docs * add error message * clean docs * comprise genot tests * change reference for GENOT * add missing docstring * Modify behaviour of `ConditionalLoader` * Update docstring * Clean GENOT docs * Improve VF * Simplify GENOT test * Better metadata wrapper in tests * Fix condition in GENOT test * Add quad cond dl * Add conf fused DL * Polish docs * Remove conditional loader * Fix link in the docs * Improve VF * Fix GENOT test * Polish docs * Remove `uniform_marginals` argument * Fix undefined variable * Update `GENOT.transport` docs * Add `diffrax` to `conf.py` * Restructure files * Fix neural init tests import * Update `docs/` * Update Monge Gap * Update MetaOT and NeuralDual * Update ICNN inits * Fix links to neural in the docs * Check for condition dim in VF * Don't use activation fn in the last layer of VF * Update assertions * Try skipping OTFM/GENOT tests temporarily * Be extra verbose when intalling packages * Remove `torch` dependency * Remove `torch` from tests in `pyproject.toml` * [ci skip] Update docstrings --------- Co-authored-by: lucaeyring <[email protected]> Co-authored-by: Michal Klein <[email protected]> Co-authored-by: nvesseron <[email protected]> Co-authored-by: Dominik Klein <[email protected]>

MUCDK added 2 commits November 22, 2023 09:54

draft of BaseSolver and UnbalancedMixin

0f91857

draft of BaseSolver and UnbalancedMixin

3706970

MUCDK added 2 commits November 22, 2023 14:12

[ci skip] continue flow matching implementation

42dd2b8

[ci skip] continue flow matching implementation

6583281

MUCDK added 6 commits November 22, 2023 16:15

[ci skip] add neural networks

f5a043c

[ci skip] add test

34bb10f

[ci skip] resolve import errors

374d051

[ci skip] MRO not working

a9e9a8c

[ci skip] basic test for flow matching passes

e4f8991

[ci skip] add tests for FM with conditions and conditional OT with FM

7869e37

MUCDK changed the title ~~draft of BaseNeuralSolver and UnbalancednessMixin~~ restructuring neural models + addition of OT-FM and GENOT Nov 23, 2023

MUCDK added 10 commits November 23, 2023 18:03

[ci skip] add genot outline

5a90dc1

[ci skip] restructure genot

c843758

[ci skip] restructure genot

ef86c54

[ci skip] fix transport

70a6173

[ci skip] flow matching tests passing

40570e6

[ci skip] add more tests genot

b0910ea

[ci skip] add more tests genot

542f512

[ci skip] add TimeSampler

c067f45

[ci skip] add docs for TimeSampler and Flow

2546afc

[ci skip] add docs for OTFlowMatching and replace jnp.ndarray by jax.…

579852f

…Array

MUCDK added 6 commits November 26, 2023 16:27

[ci skip] change init arguments of GENOT and add docstrings to GENOT

b075758

[ci skip] split nets into base_models and models

95e8707

[ci skip] add references

3b1791d

add tests for learning the rescaling factors

eca77c0

[ci skip] partially fix rescaling factor learning

62b2666

[ci skip] fix rescaling factor learning

2ceceea

michalk8 added 4 commits March 27, 2024 17:28

Fix undefined variable

ba64056

Update GENOT.transport docs

80d2924

Add diffrax to conf.py

e4aae7f

Merge remote-tracking branch 'upstream/main' into draft/neural_base_s…

128e085

…olver

JTT94 reviewed Mar 28, 2024

View reviewed changes

michalk8 added 2 commits March 29, 2024 19:05

Restructure files

1d96fac

Fix neural init tests import

ef6afd1

michalk8 added 2 commits March 29, 2024 20:12

Update docs/

73c2527

Update Monge Gap

0418e78

MUCDK commented Apr 2, 2024

View reviewed changes

michalk8 added 3 commits April 2, 2024 11:20

Update MetaOT and NeuralDual

b34b886

Update ICNN inits

67202c2

Fix links to neural in the docs

982d20b

MUCDK commented Apr 2, 2024

View reviewed changes

src/ott/neural/methods/flows/genot.py Outdated Show resolved Hide resolved

Check for condition dim in VF

7b61e05

michalk8 force-pushed the draft/neural_base_solver branch from a0fed3c to 7b61e05 Compare April 2, 2024 15:00

michalk8 added 7 commits April 2, 2024 19:28

Don't use activation fn in the last layer of VF

8819d5e

Update assertions

6f9cbcc

Try skipping OTFM/GENOT tests temporarily

9e1499b

Be extra verbose when intalling packages

b37da2a

Remove torch dependency

9c561a5

Remove torch from tests in pyproject.toml

f227d54

[ci skip] Update docstrings

6f9a77c

michalk8 self-requested a review April 3, 2024 13:00

michalk8 approved these changes Apr 3, 2024

View reviewed changes

michalk8 merged commit 713b9fc into ott-jax:main Apr 3, 2024
1 check passed

michalk8 mentioned this pull request Apr 9, 2024

New neural OT solver ENOT #503

Merged

jannisborn mentioned this pull request Apr 9, 2024

UOT-MG ExplainableML/uot-fm#2

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

restructuring neural models + addition of OT-FM and GENOT #468

restructuring neural models + addition of OT-FM and GENOT #468

MUCDK commented Nov 22, 2023 •

edited

Loading

codecov bot commented Nov 22, 2023 •

edited

Loading

MUCDK commented Nov 22, 2023

review-notebook-app bot commented Nov 26, 2023

JTT94 Mar 25, 2024

michalk8 commented Mar 29, 2024 •

edited

Loading

MUCDK Apr 2, 2024

michalk8 left a comment

		return jnp.full_like(t, fill_value=self.sigma)


		class BrownianNoiseFlow(StraightFlow):



		@pytest.fixture()
		def lin_cond_dl() -> DataLoader:

restructuring neural models + addition of OT-FM and GENOT #468

restructuring neural models + addition of OT-FM and GENOT #468

Conversation

MUCDK commented Nov 22, 2023 • edited Loading

codecov bot commented Nov 22, 2023 • edited Loading

Codecov Report

MUCDK commented Nov 22, 2023

review-notebook-app bot commented Nov 26, 2023

JTT94 Mar 25, 2024

Choose a reason for hiding this comment

michalk8 commented Mar 29, 2024 • edited Loading

MUCDK Apr 2, 2024

Choose a reason for hiding this comment

michalk8 left a comment

Choose a reason for hiding this comment

MUCDK commented Nov 22, 2023 •

edited

Loading

codecov bot commented Nov 22, 2023 •

edited

Loading

michalk8 commented Mar 29, 2024 •

edited

Loading