Core Implementation of the Simformer #1
Closed
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…-catch blocks for debugging, minor adjustments
…re-introduced some docstrings
…ous lines for ruff linting
…ion mask (no zeros, no ones)
…sion in sample batched again, everything still works
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Generally, each Masked version of other objects are provided exacly below their counterpart in the same python file, e.g.
MaskedConditionalVectorFieldEstimatoris exactly below the code block ofConditionalVectorFieldEstimatorFurthermore, regarding the automatic generation of condition masks and edge masks, (i.e., a bernoulli mask and a full ones respectively), it is only done in the
loss()function ofMaskedConditionalScoreEstimator, for the remaining methods of the Masked*-interface, they should always expect them as non-optional parameters, already in a "batched" shape (i.e., 3 dimensional input tensors: B, T, F) --- while the sameloss()function is still able to accept "un-batched" tensors (i.e., 2 dimensions inputs only: T, F), adapt them, and send them in the appropriate shape to the other methods usedFiles I modified should be the following:
sbi/analysis
sbi/analysis/plots.py # Added
.cpu()inreturn t.cpu().numpy()inensure_numpy()sbi/inference
sbi/inference/trainers/base.py# MaskedNeuralInference(ABC)sbi/inference/trainers/vfpe/base_vf_inference.py# MaskedVectorFiedlEstimatorBuilder(Protocol), MaskedVectorFieldInference(MaskedNeuralInference)sbi/inference/trainers/vfpe/simformer.py# New file, Simformer(MaskedVectorFieldInference)sbi/neural_nets
sbi/neural_nets/factory.py# simformer_nnsbi/neural_nets/estimators/base.py# MaskedConditionalEstimator, MaskedConditionalVectorFieldEstimator(MaskedConditionalEstimator)sbi/neural_nets/estimators/score_estimator.py# MaskedConditionalScoreEstimator(MaskedConditionalVectorFieldEstimator), comments just above ConditionalScoreEstimator, MaskedVEScoreEstimator(MaskedConditionalScoreEstimator)sbi/neural_nets/net_builders/vector_field_nets.py# build_vector_field_estimator accepts simformer and masked-score, MaskedSimformerBlock (rename! in MaskedTimeAdditiveBlock), MaskedDiTBlock, SimformerNet(MaskedVectorFieldNet), build_simformer_network(choose default parameters here)sbi/utils
sbi/utils/vector_field_utils.py# MaskedVectorFieldNetThe remaining files listed here can be ignored as they were my attempt to develop a joint-distribution object/pipeline (I will properly ignore them when I will apply the official commits for the future PR)
sbi/inference/joints/*# New filessbi/inference/potentials/base_potential.py# Defined MaskedBasePotential, CustomMaskedPotential(Protocol), CustomMaskedPotentialWrapper(MaskedBasePotential)sbi/inference/potentials/joint_based_potential.py# Same as joint distribution abovesbi/inference/potentials/vector_field_potential.py# Defined masked_vector_field_estimator_based_potential and MaskedVectorFieldBasedPotentialsbi/samplers/ode_solvers/base.py# MaskedNeuralODEFunc(Protocol), MaskedNeuralODE(LazyDistribution), MaskedZukoNeuralODE(MaskedNeuralODE), comment in Diffusersbi/utils/torchutils.py# ensure_latent_batch