First attempt to add conditional NFs to jester infrastructure#125
Open
haukekoehn wants to merge 7 commits into
Open
First attempt to add conditional NFs to jester infrastructure#125haukekoehn wants to merge 7 commits into
haukekoehn wants to merge 7 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds conditional-flow support: exports Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant CF as ConditionalFlow
participant Flow as UnderlyingFlow
participant JC as JacobianCorrection
User->>CF: log_prob(x, y)
CF->>CF: standardize_input(x) -> x_std
CF->>CF: standardize_cond_data(y) -> y_std
CF->>Flow: flow.log_prob(x_std, y_std)
Flow-->>CF: log_prob_std
CF->>JC: apply correction using cond stats (std or range)
JC-->>CF: log_prob_corrected
CF-->>User: return log_prob_corrected
User->>CF: sample(key, shape, y)
CF->>CF: standardize_cond_data(y) -> y_std
CF->>Flow: flow.sample(key, shape, condition=y_std)
Flow-->>CF: samples_std
CF->>CF: destandardize_output(samples_std) -> samples
CF-->>User: return samples
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Collaborator
Author
|
The code runs, but it's not thoroughly tested yet. I will experiment a little more with it and then we can see how to proceed. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (1)
jesterTOV/inference/flows/flow.py (1)
361-364: Render the new conditional-density equations as Sphinx math.The plain-text
log p(x|y)formulas will not format consistently in generated docs.As per coding guidelines, "All mathematical expressions in docstrings must use Sphinx/reStructuredText formatting with
:math:role for inline math and.. math::directive for display equations."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/flow.py` around lines 361 - 364, Update the docstring describing the Jacobian correction in flow.py to use Sphinx math formatting: replace plain-text formulas like "log p(x|y) = log p(x_std|y_std) - sum(log(std))" with inline :math: roles and display equations using the .. math:: directive; specifically render the three cases (Z-score, Min-max, None) as proper Sphinx math blocks so they appear correctly in generated docs and keep the surrounding explanatory text intact (look for the Jacobian correction section in the docstring near the comment block in flow.py).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@jesterTOV/inference/flows/config.py`:
- Line 91: Validate conditional config fields immediately when constructing the
configuration: in the config class where cond_dim and cond_parameter_names are
declared, add a sanity check that raises ValueError for invalid combinations —
require cond_dim to be None or an int >= 1 (reject 0 and negatives), require
cond_parameter_names to be either None or a list[str], forbid
cond_parameter_names when cond_dim is None or <=0, and require
len(cond_parameter_names) == cond_dim when both are provided; remove reliance on
the weaker presence-only check in train_flow.py and let train_flow.py assume
config validation has already enforced consistency.
In `@jesterTOV/inference/flows/flow.py`:
- Around line 414-416: The code calls self.flow.sample with the raw conditional
y but training standardized conditionals when config.standardize is true; update
the method so that if self.config.standardize (or equivalent flag) is enabled
you first transform y using the same routine used during training (e.g., call
the existing standardization helper such as self.standardize_input(y) or
self.standardize_condition(y) — if that helper doesn't exist, implement the same
standardization logic used during fit), then pass the standardized conditional
into self.flow.sample(key, shape, condition=standardized_y), and finally keep
the existing destandardize_output(samples) step.
- Around line 314-328: Fix standardize_cond_data in ConditionalFlow: for the
"zscore" branch divide by self.cond_data_std (not cond_data_mean); for the
"minmax" branch guard against zero ranges by using a safe_range =
jnp.where(self.cond_data_range == 0, 1, self.cond_data_range) and divide by
that; and add an explicit "none" case that returns the input data unchanged so
it doesn't reference min/max attributes that weren't initialized. Update the
standardize_cond_data method accordingly (referencing standardize_cond_data,
self.cond_data_mean, self.cond_data_std, self.cond_data_min,
self.cond_data_range).
In `@jesterTOV/inference/flows/train_flow.py`:
- Line 567: The corner-plot logic assumes `data` is an ndarray and calls
`sample()` without conditions, but earlier `data = (data, cond_samples)` for
conditional runs; update the plotting branch in the training flow to detect when
`data` is a tuple (e.g. check `isinstance(data, tuple)` or presence of
`cond_samples`) and either (a) derive an unconditional array for plotting or (b)
skip/disable the corner plot when conditions are present, or if plotting is
desired call the model's `sample(condition=...)`/`sample(cond_samples)` variant
so sampling respects the condition; ensure the check prevents the current
fallback warning by handling `cond_samples` explicitly and referencing the
`data` tuple, `cond_samples`, and the `sample()` call sites in the flow.
- Around line 651-658: The code reads cond_data_statistics unconditionally when
populating metadata in train_flow.py, which causes an UnboundLocalError when
cond_dim is None; update the block that adds conditional data stats to first
check that cond_data_statistics (or cond_dim) is defined/True before accessing
it—e.g., wrap the existing if config.standardization_method == "zscore": ...
else: ... inside a guard like if cond_data_statistics is not None (or if
cond_dim is not None and cond_data_statistics) so metadata population only runs
when cond_data_statistics exists (references: cond_data_statistics,
config.standardization_method, metadata).
---
Nitpick comments:
In `@jesterTOV/inference/flows/flow.py`:
- Around line 361-364: Update the docstring describing the Jacobian correction
in flow.py to use Sphinx math formatting: replace plain-text formulas like "log
p(x|y) = log p(x_std|y_std) - sum(log(std))" with inline :math: roles and
display equations using the .. math:: directive; specifically render the three
cases (Z-score, Min-max, None) as proper Sphinx math blocks so they appear
correctly in generated docs and keep the surrounding explanatory text intact
(look for the Jacobian correction section in the docstring near the comment
block in flow.py).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9d16684a-19f2-4060-9497-f45ddfd31655
📒 Files selected for processing (4)
jesterTOV/inference/flows/__init__.pyjesterTOV/inference/flows/config.pyjesterTOV/inference/flows/flow.pyjesterTOV/inference/flows/train_flow.py
| flow_layers: int = 1 | ||
| invert: bool = True | ||
| cond_dim: int | None = None | ||
| cond_parameter_names: list[str] | None = None |
There was a problem hiding this comment.
Fail fast on inconsistent conditional configs.
This addition exposes cond_parameter_names, but the model still accepts cond_dim=0, negative cond_dim, cond_parameter_names without cond_dim, and lists whose length does not match cond_dim. The runtime guard in jesterTOV/inference/flows/train_flow.py only checks presence, so some bad configs are silently ignored while others fail much later inside flow creation/training.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/flows/config.py` at line 91, Validate conditional config
fields immediately when constructing the configuration: in the config class
where cond_dim and cond_parameter_names are declared, add a sanity check that
raises ValueError for invalid combinations — require cond_dim to be None or an
int >= 1 (reject 0 and negatives), require cond_parameter_names to be either
None or a list[str], forbid cond_parameter_names when cond_dim is None or <=0,
and require len(cond_parameter_names) == cond_dim when both are provided; remove
reliance on the weaker presence-only check in train_flow.py and let
train_flow.py assume config validation has already enforced consistency.
Comment on lines
+314
to
+328
| if self.standardization_method=="zscore": | ||
| self.cond_data_mean = jnp.array(metadata["cond_data_mean"]) | ||
| self.cond_data_std = jnp.array(metadata["cond_data_std"]) | ||
|
|
||
| elif self.standardization_method=="minmax": | ||
| self.cond_data_min = jnp.array(metadata["cond_data_min"]) | ||
| self.cond_data_max = jnp.array(metadata["cond_data_max"]) | ||
| self.cond_data_range = self.cond_data_max - self.cond_data_min | ||
|
|
||
| def standardize_cond_data(self, data: Array) -> Array: | ||
|
|
||
| if self.standardization_method == "zscore": | ||
| return (data - self.cond_data_mean) / self.cond_data_mean | ||
| else: | ||
| return (data - self.cond_data_min) / self.cond_data_range |
There was a problem hiding this comment.
Fix the conditional standardization math.
The z-score branch divides by cond_data_mean instead of cond_data_std, the "none" case falls through to min-max attributes that were never initialized, and the min-max path does not recreate the zero-range guard from training. ConditionalFlow.log_prob() will therefore be wrong for standardized models and can crash for unstandardized ones.
🛠️ Suggested fix
if self.standardization_method=="zscore":
self.cond_data_mean = jnp.array(metadata["cond_data_mean"])
self.cond_data_std = jnp.array(metadata["cond_data_std"])
+ self.cond_data_std = jnp.where(
+ self.cond_data_std == 0, 1.0, self.cond_data_std
+ )
elif self.standardization_method=="minmax":
self.cond_data_min = jnp.array(metadata["cond_data_min"])
self.cond_data_max = jnp.array(metadata["cond_data_max"])
self.cond_data_range = self.cond_data_max - self.cond_data_min
+ self.cond_data_range = jnp.where(
+ self.cond_data_range == 0, 1.0, self.cond_data_range
+ )
def standardize_cond_data(self, data: Array) -> Array:
if self.standardization_method == "zscore":
- return (data - self.cond_data_mean) / self.cond_data_mean
- else:
+ return (data - self.cond_data_mean) / self.cond_data_std
+ if self.standardization_method == "minmax":
return (data - self.cond_data_min) / self.cond_data_range
+ return data🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/flows/flow.py` around lines 314 - 328, Fix
standardize_cond_data in ConditionalFlow: for the "zscore" branch divide by
self.cond_data_std (not cond_data_mean); for the "minmax" branch guard against
zero ranges by using a safe_range = jnp.where(self.cond_data_range == 0, 1,
self.cond_data_range) and divide by that; and add an explicit "none" case that
returns the input data unchanged so it doesn't reference min/max attributes that
weren't initialized. Update the standardize_cond_data method accordingly
(referencing standardize_cond_data, self.cond_data_mean, self.cond_data_std,
self.cond_data_min, self.cond_data_range).
Comment on lines
+414
to
+416
| samples = self.flow.sample(key, shape, condition=y) | ||
|
|
||
| samples = self.destandardize_output(samples) |
There was a problem hiding this comment.
sample() is conditioning on the wrong scale.
Training standardizes conditional inputs before fitting when config.standardize is enabled, but this method passes raw y straight into the flow. With the default config, generated samples are conditioned on mismatched inputs.
🛠️ Suggested fix
- samples = self.flow.sample(key, shape, condition=y)
+ y_std = self.standardize_cond_data(y)
+ samples = self.flow.sample(key, shape, condition=y_std)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/flows/flow.py` around lines 414 - 416, The code calls
self.flow.sample with the raw conditional y but training standardized
conditionals when config.standardize is true; update the method so that if
self.config.standardize (or equivalent flag) is enabled you first transform y
using the same routine used during training (e.g., call the existing
standardization helper such as self.standardize_input(y) or
self.standardize_condition(y) — if that helper doesn't exist, implement the same
standardization logic used during fit), then pass the standardized conditional
into self.flow.sample(key, shape, condition=standardized_y), and finally keep
the existing destandardize_output(samples) step.
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
jesterTOV/inference/flows/train_flow.py (1)
687-696:⚠️ Potential issue | 🟠 MajorInverse standardization incorrectly applied to conditional columns.
When
config.cond_dimis set,flow_samples_npcontains both primary samples and appended conditional columns (line 680). The inverse standardization at lines 690-696 appliesdata_statistics(computed for primary data) to all columns, corrupting the already-original-scale conditional values.Either inverse-transform before concatenation or slice out the conditional columns before applying the inverse.
🛠️ Suggested fix (transform before concatenation)
if config.cond_dim: flow_samples = trained_flow.sample(sample_key, (cond_samples.shape[0],), condition=cond_samples) - flow_samples_np = np.hstack((flow_samples, original_data[:, -config.cond_dim:])) + flow_samples_np = np.array(flow_samples) + # Inverse transform primary samples before concatenation + if config.standardize and data_statistics is not None: + if config.standardization_method == "zscore": + flow_samples_np = inverse_standardize_data_zscore( + flow_samples_np, data_statistics + ) + else: + flow_samples_np = inverse_standardize_data_minmax( + flow_samples_np, data_statistics + ) + flow_samples_np = np.hstack((flow_samples_np, original_cond_samples)) labels = [*config.parameter_names, *config.cond_parameter_names] else: flow_samples = trained_flow.sample(sample_key, (n_plot_samples,)) flow_samples_np = np.array(flow_samples) - - # Inverse transform samples if data was standardized - if config.standardize and data_statistics is not None: - if config.standardization_method == "zscore": - flow_samples_np = inverse_standardize_data_zscore( - flow_samples_np, data_statistics - ) - else: # minmax - flow_samples_np = inverse_standardize_data_minmax( - flow_samples_np, data_statistics - ) + # Inverse transform samples if data was standardized + if config.standardize and data_statistics is not None: + if config.standardization_method == "zscore": + flow_samples_np = inverse_standardize_data_zscore( + flow_samples_np, data_statistics + ) + else: # minmax + flow_samples_np = inverse_standardize_data_minmax( + flow_samples_np, data_statistics + ) labels = parameter_names🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/train_flow.py` around lines 687 - 696, flow_samples_np is being inverse-standardized across all columns even when config.cond_dim > 0, which applies data_statistics (computed for primary features) to appended conditional columns and corrupts them; fix by either inverse-transforming the primary-sample portion before concatenating conditionals or by slicing out the last config.cond_dim columns from flow_samples_np, applying inverse_standardize_data_zscore / inverse_standardize_data_minmax to only the primary slice (using data_statistics), and then re-attaching the conditional columns; locate this logic around flow_samples_np, config.cond_dim, inverse_standardize_data_zscore, inverse_standardize_data_minmax and data_statistics and ensure conditional columns are excluded from the inverse transform.
♻️ Duplicate comments (4)
jesterTOV/inference/flows/flow.py (3)
323-328:⚠️ Potential issue | 🔴 Critical
standardize_cond_datacrashes whenstandardization_method="none".When standardization is disabled, the
elsebranch tries to accesscond_data_minandcond_data_range, which are never initialized for the "none" case. This will raiseAttributeErrorat runtime.🛠️ Suggested fix
def standardize_cond_data(self, data: Array) -> Array: if self.standardization_method == "zscore": return (data - self.cond_data_mean) / self.cond_data_std - else: + elif self.standardization_method == "minmax": return (data - self.cond_data_min) / self.cond_data_range + else: + # No standardization - identity transform + return data🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/flow.py` around lines 323 - 328, The method standardize_cond_data currently assumes either "zscore" or min-max and will access self.cond_data_min/self.cond_data_range even when standardization_method=="none"; update standardize_cond_data to explicitly handle the "none" case (or unknown methods) by returning data unchanged (or raising a clear ValueError for unsupported methods), and avoid accessing cond_data_min/cond_data_range unless they are initialized; reference the method name standardize_cond_data and the attributes cond_data_mean, cond_data_std, cond_data_min, cond_data_range when making the change.
414-415:⚠️ Potential issue | 🔴 Critical
y_stdreferenced before assignment; rawypassed to flow.Line 414 uses
y_stdon the right-hand side before it is defined, causing aNameError. Additionally, line 415 passes the rawytoself.flow.sampleinstead of the standardized version, producing samples conditioned on the wrong scale.🛠️ Suggested fix
- y_std = self.standardize_cond_data(y_std) - samples = self.flow.sample(key, shape, condition=y) + y_std = self.standardize_cond_data(y) + samples = self.flow.sample(key, shape, condition=y_std)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/flow.py` around lines 414 - 415, The code references y_std before it is created and passes raw y into the sampler; fix by first assigning y_std = self.standardize_cond_data(y) and then call self.flow.sample(key, shape, condition=y_std) so the sampler receives the standardized conditioning; look for usages in the method containing standardize_cond_data and self.flow.sample to update the order and the condition argument (y -> y_std).
314-322:⚠️ Potential issue | 🟠 MajorMissing zero-division guards for conditional statistics.
The z-score and min-max branches load conditional statistics but do not apply the same zero-division protection that the parent
Flow.__init__applies to primary data statistics. If any conditional feature has zero variance (z-score) or zero range (min-max),standardize_cond_datawill produceinf/nan.🛠️ Suggested fix
if self.standardization_method=="zscore": self.cond_data_mean = jnp.array(metadata["cond_data_mean"]) self.cond_data_std = jnp.array(metadata["cond_data_std"]) + # Avoid division by zero + self.cond_data_std = jnp.where(self.cond_data_std == 0, 1.0, self.cond_data_std) elif self.standardization_method=="minmax": self.cond_data_min = jnp.array(metadata["cond_data_min"]) self.cond_data_max = jnp.array(metadata["cond_data_max"]) self.cond_data_range = self.cond_data_max - self.cond_data_min + # Avoid division by zero + self.cond_data_range = jnp.where(self.cond_data_range == 0, 1.0, self.cond_data_range)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/flow.py` around lines 314 - 322, The conditional-statistics loading branches for standardization_method ("zscore" and "minmax") can produce divide-by-zero when a conditional feature has zero variance/range; update the zscore branch to set self.cond_data_std = jnp.where(self.cond_data_std == 0, 1, self.cond_data_std) (or add a small epsilon) after loading metadata["cond_data_std"], and update the minmax branch to compute self.cond_data_range = self.cond_data_max - self.cond_data_min and then set self.cond_data_range = jnp.where(self.cond_data_range == 0, 1, self.cond_data_range) (or add epsilon) so that standardize_cond_data uses safe denominators; refer to standardization_method, cond_data_std, cond_data_range, and standardize_cond_data when making the change.jesterTOV/inference/flows/train_flow.py (1)
652-659:⚠️ Potential issue | 🔴 Critical
cond_data_statisticsmay beNonewhen accessed.When
config.cond_dimis set butconfig.standardizeisFalse,cond_data_statisticsremainsNone(initialized at line 560). The code then tries to accesscond_data_statistics["mean"]orcond_data_statistics["min"], raising aTypeError.🛠️ Suggested fix
# Add conditional data statistics to metadata if standardization was used - if config.cond_dim: + if config.cond_dim and cond_data_statistics is not None: if config.standardization_method == "zscore": metadata["cond_data_mean"] = cond_data_statistics["mean"].tolist() metadata["cond_data_std"] = cond_data_statistics["std"].tolist() else: metadata["cond_data_min"] = cond_data_statistics["min"].tolist() metadata["cond_data_max"] = cond_data_statistics["max"].tolist()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/train_flow.py` around lines 652 - 659, The code reads cond_data_statistics when config.cond_dim is true but cond_data_statistics can be None; update the block that adds conditional stats to metadata to first ensure cond_data_statistics is not None (or compute it from cond_data if available) before accessing keys. Specifically, in the block guarded by config.cond_dim and config.standardization_method, check cond_data_statistics is truthy and raise a clear error or skip adding stats if it's None; reference the cond_data_statistics variable, config.standardization_method, and the metadata dict when making the change so you don't accidentally access ["mean"/"std"/"min"/"max"] on None.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@jesterTOV/inference/flows/flow.py`:
- Around line 331-334: The docstring for the classmethod from_directory in
ConditionalFlow has a typo ("conditinal"); update the docstring text to read
"Load a trained conditional flow from a directory." so the wording is correct
and consistent with the ConditionalFlow class name and purpose.
- Around line 374-377: The docstring example for flow.log_prob is missing the
conditional argument y; update the example to call flow.log_prob(x, y) (i.e.,
replace flow.log_prob(data) with something like flow.log_prob(x, y) where x is
the input array and y is the corresponding conditional array), ensuring the
example uses the correct variable names from the method signature (log_prob(x,
y)) and matching shapes/types for x and y in the illustrated call.
In `@jesterTOV/inference/flows/train_flow.py`:
- Around line 198-207: The logging message "Data bounds saved for inverse
transformation" is inside the for-loop that iterates over parameter_names,
causing it to be emitted once per parameter; move that logger.info call out of
the loop so it runs once after the loop that logs per-parameter ranges in the
minmax branch (the block using standardize_data_minmax, data, data_statistics,
parameter_names, and logger). Ensure the loop only logs the per-parameter ranges
and then, immediately after the loop completes, call logger.info("Data bounds
saved for inverse transformation") once.
- Around line 679-680: The conditional call to trained_flow.sample is missing
the required shape argument; update the call to trained_flow.sample(sample_key,
shape, condition=cond_samples) by passing the desired sample shape (e.g.,
matching the number of samples and latent/event dim) — derive shape from
existing variables like cond_samples.shape[0] (batch size) and the expected
output dimensionality (use config or the flow's output dim) so that flow_samples
has the correct shape before hstack with original_data[:, -config.cond_dim:].
---
Outside diff comments:
In `@jesterTOV/inference/flows/train_flow.py`:
- Around line 687-696: flow_samples_np is being inverse-standardized across all
columns even when config.cond_dim > 0, which applies data_statistics (computed
for primary features) to appended conditional columns and corrupts them; fix by
either inverse-transforming the primary-sample portion before concatenating
conditionals or by slicing out the last config.cond_dim columns from
flow_samples_np, applying inverse_standardize_data_zscore /
inverse_standardize_data_minmax to only the primary slice (using
data_statistics), and then re-attaching the conditional columns; locate this
logic around flow_samples_np, config.cond_dim, inverse_standardize_data_zscore,
inverse_standardize_data_minmax and data_statistics and ensure conditional
columns are excluded from the inverse transform.
---
Duplicate comments:
In `@jesterTOV/inference/flows/flow.py`:
- Around line 323-328: The method standardize_cond_data currently assumes either
"zscore" or min-max and will access self.cond_data_min/self.cond_data_range even
when standardization_method=="none"; update standardize_cond_data to explicitly
handle the "none" case (or unknown methods) by returning data unchanged (or
raising a clear ValueError for unsupported methods), and avoid accessing
cond_data_min/cond_data_range unless they are initialized; reference the method
name standardize_cond_data and the attributes cond_data_mean, cond_data_std,
cond_data_min, cond_data_range when making the change.
- Around line 414-415: The code references y_std before it is created and passes
raw y into the sampler; fix by first assigning y_std =
self.standardize_cond_data(y) and then call self.flow.sample(key, shape,
condition=y_std) so the sampler receives the standardized conditioning; look for
usages in the method containing standardize_cond_data and self.flow.sample to
update the order and the condition argument (y -> y_std).
- Around line 314-322: The conditional-statistics loading branches for
standardization_method ("zscore" and "minmax") can produce divide-by-zero when a
conditional feature has zero variance/range; update the zscore branch to set
self.cond_data_std = jnp.where(self.cond_data_std == 0, 1, self.cond_data_std)
(or add a small epsilon) after loading metadata["cond_data_std"], and update the
minmax branch to compute self.cond_data_range = self.cond_data_max -
self.cond_data_min and then set self.cond_data_range =
jnp.where(self.cond_data_range == 0, 1, self.cond_data_range) (or add epsilon)
so that standardize_cond_data uses safe denominators; refer to
standardization_method, cond_data_std, cond_data_range, and
standardize_cond_data when making the change.
In `@jesterTOV/inference/flows/train_flow.py`:
- Around line 652-659: The code reads cond_data_statistics when config.cond_dim
is true but cond_data_statistics can be None; update the block that adds
conditional stats to metadata to first ensure cond_data_statistics is not None
(or compute it from cond_data if available) before accessing keys. Specifically,
in the block guarded by config.cond_dim and config.standardization_method, check
cond_data_statistics is truthy and raise a clear error or skip adding stats if
it's None; reference the cond_data_statistics variable,
config.standardization_method, and the metadata dict when making the change so
you don't accidentally access ["mean"/"std"/"min"/"max"] on None.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 2815381c-fe1c-4f7b-acf0-afa64357dbe4
📒 Files selected for processing (2)
jesterTOV/inference/flows/flow.pyjesterTOV/inference/flows/train_flow.py
Comment on lines
+198
to
+207
| else: # minmax | ||
| logger.info("Standardizing data using min-max [0, 1] scaling...") | ||
| data, data_statistics = standardize_data_minmax(data) | ||
| logger.info("Standardized data ranges:") | ||
| for i, name in enumerate(parameter_names): | ||
| logger.info( | ||
| f" {name}: [{data[:, i].min():.3f}, {data[:, i].max():.3f}]" | ||
| ) | ||
| logger.info("Data bounds saved for inverse transformation") | ||
|
|
There was a problem hiding this comment.
Logging statement inside loop causes repeated output.
Line 206 is indented inside the for loop, so "Data bounds saved for inverse transformation" will be logged once per parameter instead of once total. Compare with the z-score branch where the equivalent message (line 197) is correctly outside the loop.
🛠️ Suggested fix
else: # minmax
logger.info("Standardizing data using min-max [0, 1] scaling...")
data, data_statistics = standardize_data_minmax(data)
logger.info("Standardized data ranges:")
for i, name in enumerate(parameter_names):
logger.info(
f" {name}: [{data[:, i].min():.3f}, {data[:, i].max():.3f}]"
)
- logger.info("Data bounds saved for inverse transformation")
+ logger.info("Data bounds saved for inverse transformation")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/flows/train_flow.py` around lines 198 - 207, The logging
message "Data bounds saved for inverse transformation" is inside the for-loop
that iterates over parameter_names, causing it to be emitted once per parameter;
move that logger.info call out of the loop so it runs once after the loop that
logs per-parameter ranges in the minmax branch (the block using
standardize_data_minmax, data, data_statistics, parameter_names, and logger).
Ensure the loop only logs the per-parameter ranges and then, immediately after
the loop completes, call logger.info("Data bounds saved for inverse
transformation") once.
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (3)
jesterTOV/inference/flows/flow.py (2)
416-417:⚠️ Potential issue | 🔴 CriticalUse the standardized condition in
sample().Line 416 reads
y_stdbefore it exists, and Line 417 still sends rawyinto the flow. So this method currently fails immediately, and it would still condition on the wrong scale after fixing theNameError.Suggested fix
- y_std = self.standardize_cond_data(y_std) - samples = self.flow.sample(key, shape, condition=y) + y_std = self.standardize_cond_data(y) + samples = self.flow.sample(key, shape, condition=y_std)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/flow.py` around lines 416 - 417, The code calls self.standardize_cond_data but then passes the unstandardized y into self.flow.sample and also reads y_std before it's defined; fix by assigning the result of self.standardize_cond_data(y) to y_std (or rename the variable consistently) and then pass that standardized variable into self.flow.sample as the condition argument; update the block around standardize_cond_data and flow.sample in the same method so sample(...) receives y_std instead of y.
314-328:⚠️ Potential issue | 🔴 CriticalHandle the non-standardized conditional path explicitly.
standardize_cond_data()still assumes min/max stats exist whenever the method is not"zscore". Withstandardize=False,ConditionalFlownever initializescond_data_min/cond_data_range, so the first conditionallog_prob()orsample()call will crash. The min-max path also rebuildscond_data_rangewithout the zero-range guard used during training, which can reintroduceinf/nanfor constant conditional features.Suggested fix
if self.standardization_method=="zscore": self.cond_data_mean = jnp.array(metadata["cond_data_mean"]) self.cond_data_std = jnp.array(metadata["cond_data_std"]) + self.cond_data_std = jnp.where(self.cond_data_std == 0, 1.0, self.cond_data_std) elif self.standardization_method=="minmax": self.cond_data_min = jnp.array(metadata["cond_data_min"]) self.cond_data_max = jnp.array(metadata["cond_data_max"]) self.cond_data_range = self.cond_data_max - self.cond_data_min + self.cond_data_range = jnp.where( + self.cond_data_range == 0, 1.0, self.cond_data_range + ) def standardize_cond_data(self, data: Array) -> Array: if self.standardization_method == "zscore": return (data - self.cond_data_mean) / self.cond_data_std - else: + if self.standardization_method == "minmax": return (data - self.cond_data_min) / self.cond_data_range + return data🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/flow.py` around lines 314 - 328, standardize_cond_data currently assumes a minmax path when standardization_method != "zscore", causing crashes if ConditionalFlow was created with standardize=False and also reintroducing divide-by-zero for constant features; update the metadata-loading and standardize_cond_data logic to explicitly handle a non-standardized case (e.g., standardization_method is None or "none") by returning data unchanged in standardize_cond_data, and for the "minmax" branch ensure cond_data_min/cond_data_max/cond_data_range are initialized from metadata (or computed) and clamp cond_data_range with a small eps (or replace zeros with 1) to avoid division by zero—refer to symbols standardization_method, standardize_cond_data, cond_data_min, cond_data_max, cond_data_range, cond_data_mean, cond_data_std to locate and modify the code.jesterTOV/inference/flows/train_flow.py (1)
652-659:⚠️ Potential issue | 🔴 CriticalOnly persist conditional stats when they were actually computed.
This block runs for every conditional training job, but
cond_data_statisticsis only populated whenconfig.standardizeis true. Withcond_dimset and standardization disabled, saving the model crashes here oncond_data_statistics["..."].Suggested fix
- if config.cond_dim: + if config.cond_dim and config.standardize and cond_data_statistics is not None: if config.standardization_method == "zscore": metadata["cond_data_mean"] = cond_data_statistics["mean"].tolist() metadata["cond_data_std"] = cond_data_statistics["std"].tolist() else: metadata["cond_data_min"] = cond_data_statistics["min"].tolist() metadata["cond_data_max"] = cond_data_statistics["max"].tolist()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/train_flow.py` around lines 652 - 659, The code writes conditional stats unconditionally when config.cond_dim is set but cond_data_statistics is only created when standardization was enabled; guard these writes by checking the standardization flag and existence of cond_data_statistics (e.g., only proceed if config.standardize is True and cond_data_statistics is not None), then branch on config.standardization_method ("zscore") to persist mean/std or min/max into metadata using the existing keys (cond_data_mean, cond_data_std, cond_data_min, cond_data_max) so you never index cond_data_statistics when it wasn't computed.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@jesterTOV/inference/flows/flow.py`:
- Around line 354-365: Update the ConditionalFlow.log_prob() docstring to use a
raw string (r""") and convert the probability change-of-variable lines into
Sphinx math: use :math: for inline symbols and a ``.. math::`` block for the
display equations showing Z-score, Min-max, and None corrections (e.g. .. math::
\\log p(x\\mid y) = \\log p(x_{std}\\mid y_{std}) - \\sum \\log(\\mathrm{std}),
etc.), ensuring each formula is valid LaTeX inside the math directive and the
prose references those math expressions via :math: where needed.
In `@jesterTOV/inference/flows/train_flow.py`:
- Around line 679-694: The sampled primary parameters are being flattened with
reshape(-1, 1) which collapses multi-dimensional samples and breaks alignment
with labels; instead keep flow_samples_np as an (n_samples, dim) matrix (i.e.,
do not force a single-column reshape) and horizontally concatenate it with the
matching conditional rows from original_data (use original_data[:
flow_samples_np.shape[0], -config.cond_dim:] to ensure row counts match). Update
the hstack call that builds flow_samples_np and ensure labels =
[*config.parameter_names, *config.cond_parameter_names] still matches the
resulting column count.
- Around line 545-550: The check using len(config.cond_parameter_names) can
raise TypeError if cond_parameter_names is None; update the validation in
train_flow.py (around the block that references
FlowTrainingConfig.cond_parameter_names and config.cond_dim) to first guard that
config.cond_parameter_names is not None/empty and if config.cond_dim > 0 then
raise a ValueError when names are missing, and otherwise if it is present
validate len(config.cond_parameter_names) == config.cond_dim and raise the
existing ValueError with the same message; ensure you reference
FlowTrainingConfig.cond_parameter_names and config.cond_dim in the new guard
logic.
---
Duplicate comments:
In `@jesterTOV/inference/flows/flow.py`:
- Around line 416-417: The code calls self.standardize_cond_data but then passes
the unstandardized y into self.flow.sample and also reads y_std before it's
defined; fix by assigning the result of self.standardize_cond_data(y) to y_std
(or rename the variable consistently) and then pass that standardized variable
into self.flow.sample as the condition argument; update the block around
standardize_cond_data and flow.sample in the same method so sample(...) receives
y_std instead of y.
- Around line 314-328: standardize_cond_data currently assumes a minmax path
when standardization_method != "zscore", causing crashes if ConditionalFlow was
created with standardize=False and also reintroducing divide-by-zero for
constant features; update the metadata-loading and standardize_cond_data logic
to explicitly handle a non-standardized case (e.g., standardization_method is
None or "none") by returning data unchanged in standardize_cond_data, and for
the "minmax" branch ensure cond_data_min/cond_data_max/cond_data_range are
initialized from metadata (or computed) and clamp cond_data_range with a small
eps (or replace zeros with 1) to avoid division by zero—refer to symbols
standardization_method, standardize_cond_data, cond_data_min, cond_data_max,
cond_data_range, cond_data_mean, cond_data_std to locate and modify the code.
In `@jesterTOV/inference/flows/train_flow.py`:
- Around line 652-659: The code writes conditional stats unconditionally when
config.cond_dim is set but cond_data_statistics is only created when
standardization was enabled; guard these writes by checking the standardization
flag and existence of cond_data_statistics (e.g., only proceed if
config.standardize is True and cond_data_statistics is not None), then branch on
config.standardization_method ("zscore") to persist mean/std or min/max into
metadata using the existing keys (cond_data_mean, cond_data_std, cond_data_min,
cond_data_max) so you never index cond_data_statistics when it wasn't computed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 42186b54-53b8-4748-b4df-e16239bd559c
📒 Files selected for processing (2)
jesterTOV/inference/flows/flow.pyjesterTOV/inference/flows/train_flow.py
Comment on lines
+354
to
+365
| """ | ||
| Evaluate log probability of samples under the flow. | ||
|
|
||
| If standardization was used, input data is automatically standardized | ||
| before evaluation and Jacobian correction is applied. If not, operations | ||
| are identity (no-op). | ||
|
|
||
| The Jacobian correction accounts for the change of variables: | ||
| - Z-score: log p(x|y) = log p(x_std|y_std) - sum(log(std)) | ||
| - Min-max: log p(x|y) = log p(x_std|y_std) - sum(log(max - min)) | ||
| - None: log p(x|y) = log p(x_std|y_std) (no correction) | ||
|
|
There was a problem hiding this comment.
Format the new probability equations as Sphinx math.
The ConditionalFlow.log_prob() docstring adds math-heavy text, but it is plain prose right now. Please switch it to Sphinx math markup so the docs render consistently.
As per coding guidelines "All mathematical expressions in docstrings must use Sphinx/reStructuredText formatting with :math: role for inline math and .. math:: directive for display equations" and "Use raw strings (r""") for docstrings containing LaTeX to avoid Python escape sequence warnings".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/flows/flow.py` around lines 354 - 365, Update the
ConditionalFlow.log_prob() docstring to use a raw string (r""") and convert the
probability change-of-variable lines into Sphinx math: use :math: for inline
symbols and a ``.. math::`` block for the display equations showing Z-score,
Min-max, and None corrections (e.g. .. math:: \\log p(x\\mid y) = \\log
p(x_{std}\\mid y_{std}) - \\sum \\log(\\mathrm{std}), etc.), ensuring each
formula is valid LaTeX inside the math directive and the prose references those
math expressions via :math: where needed.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
jesterTOV/inference/flows/train_flow.py (1)
170-208:⚠️ Potential issue | 🟡 MinorMutable default argument and logging placement issues.
Mutable default argument: Using
[]as default forparameter_namesis a Python anti-pattern that can cause unexpected behavior if the list is mutated. UseNoneinstead.Logging inside loop (Line 206): The "Data bounds saved" message is inside the
forloop, causing it to log once per parameter instead of once total.🛠️ Suggested fix
def standardize_data( data: np.ndarray, standardization_method: str, - parameter_names: list[str] = [] + parameter_names: list[str] | None = None ) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: """ Standardize data based on the selected standardiation method. Args: data: Array of shape (n_samples, n_features) standardization_method: str Which method to standardize with. Can either be 'zscore' for zero mean and unit std, otherwise min-max scaling will be applied. parameter_names: Parameter names for which the rescaled data range should be printed. Defaults to []. + + Returns: + standardized_data: Data standardized according to method + data_statistics: Dictionary with standardization parameters """ + if parameter_names is None: + parameter_names = [] if standardization_method == "zscore": logger.info("Standardizing data using z-score (mean=0, std=1)...") data, data_statistics = standardize_data_zscore(data) logger.info("Standardized data statistics:") for i, name in enumerate(parameter_names): logger.info( f" {name}: mean={data[:, i].mean():.3f}, std={data[:, i].std():.3f}" ) logger.info("Data mean and std saved for inverse transformation") else: # minmax logger.info("Standardizing data using min-max [0, 1] scaling...") data, data_statistics = standardize_data_minmax(data) logger.info("Standardized data ranges:") for i, name in enumerate(parameter_names): logger.info( f" {name}: [{data[:, i].min():.3f}, {data[:, i].max():.3f}]" ) - logger.info("Data bounds saved for inverse transformation") + logger.info("Data bounds saved for inverse transformation") return data, data_statistics🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/train_flow.py` around lines 170 - 208, The function standardize_data uses a mutable default for parameter_names and logs the "Data bounds saved for inverse transformation" inside the per-parameter loop; change the signature to use parameter_names: Optional[list[str]] = None and at the start of standardize_data set parameter_names = parameter_names or [] so callers get a fresh list, and move the logger.info("Data bounds saved for inverse transformation") call out of the for i, name in enumerate(parameter_names) loop (place it after the loop, analogous to the zscore branch) so it logs once per operation; keep references to data_statistics, standardize_data_zscore, and standardize_data_minmax intact.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@jesterTOV/inference/flows/train_flow.py`:
- Around line 652-660: The metadata population block assumes
cond_data_statistics is set; guard against it by only reading
cond_data_statistics when it is not None (or when config.standardize is True)
before accessing keys like cond_data_statistics["mean"]/["std"]/["min"]/["max"];
update the conditional around config.cond_dim to also require
cond_data_statistics is not None (or config.standardize) and only then assign
metadata["cond_data_mean"/"cond_data_std"/"cond_data_min"/"cond_data_max"],
otherwise skip or set safe defaults—this change should be applied to the block
that references config.cond_dim and cond_data_statistics in the train_flow.py
flow.
- Around line 680-697: The conditional sampling uses a fixed batch size of (1,)
which mismatches cond_samples' batch size and causes the hstack with
original_data to fail; change the trained_flow.sample call to generate the same
number of samples as cond_samples (use cond_samples.shape[0] or n_conditions) so
flow_samples_np has shape (n_conditions, n_params), then keep the existing
inverse_standardize_data_zscore/inverse_standardize_data_minmax and hstack
logic; reference trained_flow.sample, cond_samples, flow_samples_np,
inverse_standardize_data_zscore/inverse_standardize_data_minmax,
config.parameter_names, config.cond_dim, original_data, and labels to locate
where to adjust the sampling shape and subsequent reshape.
---
Duplicate comments:
In `@jesterTOV/inference/flows/train_flow.py`:
- Around line 170-208: The function standardize_data uses a mutable default for
parameter_names and logs the "Data bounds saved for inverse transformation"
inside the per-parameter loop; change the signature to use parameter_names:
Optional[list[str]] = None and at the start of standardize_data set
parameter_names = parameter_names or [] so callers get a fresh list, and move
the logger.info("Data bounds saved for inverse transformation") call out of the
for i, name in enumerate(parameter_names) loop (place it after the loop,
analogous to the zscore branch) so it logs once per operation; keep references
to data_statistics, standardize_data_zscore, and standardize_data_minmax intact.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 22cf5c49-35be-4110-9939-23e133b0adaf
📒 Files selected for processing (1)
jesterTOV/inference/flows/train_flow.py
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
jesterTOV/inference/flows/train_flow.py (1)
199-207:⚠️ Potential issue | 🟡 MinorMove the min-max “bounds saved” log outside the loop.
On Line 206,
"Data bounds saved for inverse transformation"is logged once per parameter instead of once per standardization step.💡 Suggested change
else: # minmax logger.info("Standardizing data using min-max [0, 1] scaling...") data, data_statistics = standardize_data_minmax(data) logger.info("Standardized data ranges:") for i, name in enumerate(parameter_names): logger.info( f" {name}: [{data[:, i].min():.3f}, {data[:, i].max():.3f}]" ) - logger.info("Data bounds saved for inverse transformation") + logger.info("Data bounds saved for inverse transformation")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/train_flow.py` around lines 199 - 207, The log "Data bounds saved for inverse transformation" is inside the loop that iterates parameter_names and thus gets emitted once per parameter; move that logger.info call out of the for-loop in train_flow.py so it runs once after calling standardize_data_minmax (after the loop that logs each parameter's min/max). Update the block around logger.info("Standardizing data using min-max [0, 1] scaling..."), the call to standardize_data_minmax(data), and the for i, name in enumerate(parameter_names): loop to place the single "Data bounds saved..." message immediately after the loop (or immediately after data_statistics is set) so it logs once per standardization step.
🧹 Nitpick comments (1)
jesterTOV/inference/flows/train_flow.py (1)
678-697: Applyn_plot_samplescap in the conditional plotting path too.
n_plot_samplesis computed on Line 678 but ignored whenconfig.cond_dimis set; the conditional branch always uses all conditions. Capping here keeps plot generation bounded and consistent with the non-conditional path.💡 Suggested change
if config.cond_dim: # get flow samples and untransform them - n_cond = cond_samples.shape[0] - flow_samples = trained_flow.sample(sample_key, (n_cond,), condition=cond_samples) # generate one sample per condition point + n_cond = min(n_plot_samples, cond_samples.shape[0]) + cond_samples_plot = cond_samples[:n_cond] + flow_samples = trained_flow.sample( + sample_key, (n_cond,), condition=cond_samples_plot + ) # generate one sample per condition point flow_samples_np = np.array(flow_samples) @@ flow_samples_np = np.hstack( (flow_samples_np.reshape(-1, len(config.parameter_names)), original_data[:n_cond, -config.cond_dim:]) )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/flows/train_flow.py` around lines 678 - 697, The conditional plotting branch ignores the n_plot_samples cap; modify the code around the config.cond_dim path to limit n_cond to at most n_plot_samples by computing n_cond = min(cond_samples.shape[0], n_plot_samples) and slice cond_samples and original_data accordingly before calling trained_flow.sample(sample_key, (n_cond,), condition=cond_samples) and before stacking into flow_samples_np so that flow_samples_np and original_data[:n_cond, ...] use the capped size.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@jesterTOV/inference/flows/train_flow.py`:
- Around line 175-187: Update the docstring for the function that standardizes
data (references: parameter names 'standardization_method' and
'parameter_names') to use Sphinx math roles: replace inline math like "mean=0"
and "std=1" with :math:`mean = 0` and :math:`std = 1`, and replace range
notation "[0, 1]" with :math:`[0, 1]`; if any expression should be displayed on
its own line convert it to a block using the .. math:: directive. Ensure the
textual description still mentions 'zscore' and min-max scaling but uses
:math:`...` or .. math:: for all mathematical fragments.
---
Duplicate comments:
In `@jesterTOV/inference/flows/train_flow.py`:
- Around line 199-207: The log "Data bounds saved for inverse transformation" is
inside the loop that iterates parameter_names and thus gets emitted once per
parameter; move that logger.info call out of the for-loop in train_flow.py so it
runs once after calling standardize_data_minmax (after the loop that logs each
parameter's min/max). Update the block around logger.info("Standardizing data
using min-max [0, 1] scaling..."), the call to standardize_data_minmax(data),
and the for i, name in enumerate(parameter_names): loop to place the single
"Data bounds saved..." message immediately after the loop (or immediately after
data_statistics is set) so it logs once per standardization step.
---
Nitpick comments:
In `@jesterTOV/inference/flows/train_flow.py`:
- Around line 678-697: The conditional plotting branch ignores the
n_plot_samples cap; modify the code around the config.cond_dim path to limit
n_cond to at most n_plot_samples by computing n_cond =
min(cond_samples.shape[0], n_plot_samples) and slice cond_samples and
original_data accordingly before calling trained_flow.sample(sample_key,
(n_cond,), condition=cond_samples) and before stacking into flow_samples_np so
that flow_samples_np and original_data[:n_cond, ...] use the capped size.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: f7961070-4ff9-4245-b2ab-cf89f104a5e2
📒 Files selected for processing (1)
jesterTOV/inference/flows/train_flow.py
Comment on lines
+175
to
+187
| """ | ||
| Standardize data based on the selected standardiation method. | ||
|
|
||
| Args: | ||
| data: Array of shape (n_samples, n_features) | ||
| standardization_method: str | ||
| Which method to standardize with. | ||
| Can either be 'zscore' for zero mean and unit std, | ||
| otherwise min-max scaling will be applied. | ||
| parameter_names: | ||
| Parameter names for which the rescaled data range should be printed. | ||
| Defaults to []. | ||
| """ |
There was a problem hiding this comment.
Format math in docstrings using Sphinx math roles.
This new docstring includes mathematical expressions (mean=0, std=1, [0, 1]) but not Sphinx math markup.
As per coding guidelines, "All mathematical expressions in docstrings must use Sphinx/reStructuredText formatting with :math: role for inline math and .. math:: directive for display equations".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/flows/train_flow.py` around lines 175 - 187, Update the
docstring for the function that standardizes data (references: parameter names
'standardization_method' and 'parameter_names') to use Sphinx math roles:
replace inline math like "mean=0" and "std=1" with :math:`mean = 0` and
:math:`std = 1`, and replace range notation "[0, 1]" with :math:`[0, 1]`; if any
expression should be displayed on its own line convert it to a block using the
.. math:: directive. Ensure the textual description still mentions 'zscore' and
min-max scaling but uses :math:`...` or .. math:: for all mathematical
fragments.
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Introduce extra functions to accommodate conditional flows in jester.
Summary by CodeRabbit
New Features
Configuration
Bug Fixes