ENH: Reviser par Sonnet 5

[ntustison]

Revue de code — src/antsnormflows/ (hors core.py, déjà traité)

Revue de 45 fichiers (~8000 lignes) : distributions/, flows/ (+ affine/, neural_spline/), nets/, sampling/, utils/ restants, transforms.py, __init__.py.

Statut : toutes les corrections listées ci-dessous ont été appliquées, à l'exception du renommage des deux classes ConditionalDiagGaussian (choix explicite : documentation croisée seulement, pas de renommage, pour ne pas casser l'API publique et les tests). La factorisation GlowBlock2d/3d et Invertible1x1Conv/1x1x1Conv (section "Autres observations") n'a pas non plus été effectuée, comme indiqué à l'époque — trop risqué à faire sans exécution possible des tests GPU. La syntaxe de tous les fichiers modifiés a été vérifiée (ast.parse), et les tests existants ont été relus pour confirmer qu'aucun ne dépend du comportement changé.

Le problème le plus important : log_det/log_prob sans device=/dtype=

Un même bug se répète dans au moins 8 endroits : un tenseur de log-déterminant ou de log-probabilité est créé avec torch.zeros(...)/torch.ones(...) sans préciser device=. Par défaut ces tenseurs sont créés sur CPU. Si le modèle tourne sur GPU, l'opération qui combine ensuite ce tenseur avec un tenseur GPU (+=, ou pire .float() puis addition) lève soit une erreur explicite (RuntimeError: Expected all tensors to be on the same device), soit dans un cas précis (voir plus bas) une AttributeError. C'est le même bug que j'ai corrigé dans distributions/encoder.py-style ailleurs, mais il n'avait pas été traité ici.

Fichier	Ligne	Code	Impact
flows/base.py	66	total_logabsdet = torch.zeros(batch_size) dans Composite._cascade	Composite (exportée publiquement) plante sur GPU
flows/mixing.py	342	logabsdet = torch.zeros(batch_size) dans _Permutation._permute	Utilisé par LULinearPermute (exportée), plante sur GPU
flows/mixing.py	393, 412	torch.ones(outputs.shape[0]) dans le chemin caché (using_cache=True) de _Linear	Plante sur GPU seulement si le cache est activé
flows/mixing.py	609	identity = torch.eye(self.features, self.features) dans _LULinear.weight_inverse	Plante sur GPU dès qu'on appelle weight_inverse()
distributions/encoder.py	45, 49, 66, 70	torch.zeros(z.size()[0:2]) dans Dirac/Uniform (encoders VAE)	Dirac est le q0 par défaut de NormalizingFlowVAE — plante sur GPU dès la première utilisation par défaut
distributions/encoder.py	104	torch.randn((batch_size, num_samples, self.d), device=x.device) dans ConstDiagGaussian.forward	Plante avec AttributeError si x=None (cas documenté comme valide)

Recommandation : passer device= (et dtype= où pertinent) partout ci-dessus, en suivant le pattern déjà utilisé correctement ailleurs dans le même fichier (ex. distributions/base.py).

Bug transverse : log_d.float() suppose que log_det est toujours un tenseur

flows/reshape.py (Split, Merge, Squeeze2d, Squeeze3d) retourne log_det = 0 (entier Python), pas un tenseur — alors que flows/base.py définit justement un helper zero_log_det_like_z(z) pour éviter ça, mais ces classes ne l'utilisent pas.

C'est sans conséquence tant que le log_det est seulement accumulé via += (un tenseur += 0 ne pose pas de problème), ce qui est le cas dans MultiscaleFlow (qui a une garde explicite .float() if torch.is_tensor(log_det_) else log_det_) et dans AffineCouplingBlock. Mais core.py::_apply_flow_sequence (utilisée par NormalizingFlow/ConditionalNormalizingFlow) fait log_det += log_d.float() sans cette garde. Si quelqu'un construit un NormalizingFlow (pas MultiscaleFlow) avec Split, Merge, Squeeze2d ou Squeeze3d dans sa liste de flows, forward_and_log_det/inverse_and_log_det plante avec AttributeError: 'int' object has no attribute 'float'.

Deux corrections possibles : (a) faire retourner à Split/Merge/Squeeze2d/Squeeze3d un vrai tenseur via zero_log_det_like_z(z), ce qui est la correction la plus propre et cohérente avec le reste du code ; ou (b) ajouter la même garde torch.is_tensor(...) dans _apply_flow_sequence. Je recommande (a).

Bugs ponctuels

Fichier	Ligne	Problème
nets/mlp.py	52	NotImplementedError(...) construite mais jamais levée (raise manquant) — un output_fn invalide est silencieusement ignoré au lieu de lever une erreur
distributions/base.py	GaussianPCA.forward/log_prob	La boucle de retry Cholesky (for _ in range(5): try: L = ... except RuntimeError: jitter *= 10) peut se terminer sans jamais réussir à calculer L, causant un UnboundLocalError confus au lieu d'une erreur claire
distributions/encoder.py	ConstDiagGaussian.forward	Crash AttributeError si x=None, cas pourtant documenté et supporté par la signature
transforms.py	Shift.forward/inverse	z -= self.shift / z += self.shift modifient le tenseur en place, contrairement à toutes les autres flows du code (qui clonent avant de muter, ex. PeriodicWrap/PeriodicShift). Risque d'erreur autograd (RuntimeError de version counter) si z est utilisé ailleurs ou nécessite un gradient
flows/mixing.py	563	except: nu (attrape tout, y compris KeyboardInterrupt) dans _LULinear.inverse_no_cache — à remplacer par une exception précise
flows/affine/autoregressive.py	MaskedAffineAutoregressive.__init__	self.features = features est assigné avant super().__init__(made) — fonctionne par chance ici (un entier ne déclenche pas la vérification nn.Module), mais c'est un anti-pattern fragile
utils/splines.py	16-20	conditional_compile = torch.compile appliqué automatiquement à search_sorted/unconstrained_rational_quadratic_spline dès que la variable d'environnement CI n'est pas "true" — comportement de production différent du comportement en CI, jamais testé, risque de recompilations ou d'échecs silencieux liés à torch.compile en dehors de CI
utils/splines.py	185	assert (discriminant >= 0).all() — les assert sont supprimés en mode python -O; à remplacer par une vérification explicite si c'est une garde de sécurité numérique importante

Duplication de nom : deux classes ConditionalDiagGaussian

distributions/base.py::ConditionalDiagGaussian (une distribution q0, prend un context_encoder réseau de neurones) et distributions/target.py::ConditionalDiagGaussian (une distribution cible p, où context est directement [loc, scale] concaténés) portent le même nom mais n'ont rien en commun — ni la signature, ni la sémantique de context. Aucune des deux n'est exportée à la racine du package (antsnormflows.distributions.ConditionalDiagGaussian n'existe pas), les tests important explicitement from antsnormflows.distributions.base import ConditionalDiagGaussian ou from antsnormflows.distributions.target import ConditionalDiagGaussian. Risque réel de confusion / mauvais import silencieux. De même, distributions/base.py::Uniform (distribution q0 non conditionnelle) est masquée par distributions/encoder.py::Uniform (encodeur VAE conditionnel) au niveau du package : antsnormflows.distributions.Uniform résout vers la version de encoder.py, pas celle de base.py.

Recommandation : renommer l'une des deux paires de classes (ex. base.ConditionalDiagGaussian → ConditionalDiagGaussianEncoder, ou target.ConditionalDiagGaussian → ConditionalDiagGaussianTarget), et exporter explicitement (ou explicitement exclure avec un commentaire) chaque nom ambigu dans __init__.py.

Autres observations (mineures, pas bloquantes)

• distributions/linear_interpolation.py::LinearInterpolation n'hérite pas de nn.Module bien qu'il enveloppe potentiellement deux distributions qui en sont. Si dist1/dist2 sont entraînables, leurs paramètres ne seront pas suivis par .parameters()/.to(device) du module englobant.
• sampling/hais.py::HAIS.layers est une liste Python simple (pas nn.ModuleList), et HAIS n'est pas un nn.Module — cohérent avec un usage "inférence seulement", mais à signaler si un jour on veut entraîner les paramètres log_step_size/log_mass des couches HMC qu'il contient.
• __init__.py : le except Exception: autour de l'import de core.py est très large ; en cas de vraie erreur dans core.py, l'utilisateur obtient NormalizingFlow = None puis un TypeError confus plus tard, au lieu du traceback original.
• distributions/prior.py::PriorDistribution a un __init__ qui lève NotImplementedError inconditionnellement (au lieu d'utiliser abc.ABC) ; les sous-classes évitent le problème simplement en ne l'appelant jamais via super().__init__(). Fonctionne, mais fragile si une des classes venait à changer.
• utils/masks.py::create_mid_split_binary_mask et create_random_binary_mask semblent être du code mort (jamais appelés dans src/).
• distributions/base.py::GaussianMixture utilise torch.log(torch.softmax(...)) plutôt que F.log_softmax(...), moins stable numériquement.
• Duplication de code : GlowBlock2d/GlowBlock3d (flows/affine/glow.py) et Invertible1x1Conv/Invertible1x1x1Conv (flows/mixing.py) sont du code quasiment dupliqué 2D/3D — factorisable, mais risqué à toucher sans tests GPU disponibles.

Ce que je n'ai pas fait

Je n'ai pas vérifié en détail la justesse mathématique de chaque formule de densité (Sinusoidal, Smiley, TwoModes, etc. dans distributions/prior.py) ni les fichiers nets/lipschitz.py (~600 lignes, code repris de rtqichen/residual-flows) au même niveau de détail que le reste — repérage rapide de patterns à risque seulement, pas de relecture ligne à ligne complète.