Skip to content

fix(xreg): handle empty dimensions in padded JAX arrays - #478

Open
Bitshank-2338 wants to merge 2 commits into
google-research:masterfrom
Bitshank-2338:fix/xreg-empty-padding-197
Open

Bitshank-2338 wants to merge 2 commits into
google-research:masterfrom
Bitshank-2338:fix/xreg-empty-padding-197

Conversation

@Bitshank-2338

Copy link
Copy Markdown

Summary

  • preserve zero-length axes when padding XReg arrays to power-of-two dimensions
  • keep the current and archived v1 XReg implementations behaviorally aligned
  • add coverage for empty 1D/2D arrays and the single-category, no-intercept regression path

Root cause

_to_padded_jax_array called math.log2(0) for arrays with an empty dimension. A valid way to reach this state is a single-category feature encoded with drop="first" while the intercept is disabled, which produces an (n, 0) design matrix.

The padding helper now treats a zero-length dimension as requiring zero padding while continuing to power-of-two pad every non-empty dimension.

Testing

  • PYTHONPATH=src python -m pytest tests -q — 68 passed
  • PYTHONPATH=src python -m pytest tests/test_xreg_lib.py -q — 5 passed
  • python -m ruff check tests/test_xreg_lib.py
  • git diff --check

Fixes #197

@google-cla

google-cla Bot commented Sep 2, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@rajpratham1 rajpratham1 left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me. The change correctly handles zero-sized dimensions without affecting the existing padding behavior for non-empty arrays. The added regression tests cover both empty-dimension cases and the existing padding behavior, so this looks safe to merge.

@sylvesterkaczmarek sylvesterkaczmarek left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now that zero-column design matrices are supported, max_rows_per_col exposes another edge case here. When x_train.shape[1] == 0, w = ncols * max_rows_per_col becomes 0, so any positive max_rows_per_col selects an empty row set and also truncates flat_targets to length 0. The fitted predictions remain zero because there are no predictors, but debug_info=True no longer returns the documented flattened target vector. Could the subsampling branch be skipped when ncols == 0, with a regression combining the single-category/no-intercept case, max_rows_per_col > 0, and debug_info=True?

@Bitshank-2338

Copy link
Copy Markdown
Author

Thanks @sylvesterkaczmarek, I reproduced this with the single-category/no-intercept case and debug_info=True. With max_rows_per_col=0, the returned target vector is [1, 2, 3, 0] (including padding) and the training matrix has shape (4, 0). Setting max_rows_per_col to either 1 or 10 returns an empty target vector and a (0, 0) training matrix, while the predictions remain zero.

Your diagnosis is correct: ncols * max_rows_per_col becomes zero and the subsampling branch removes every row. Skipping that branch when ncols == 0 makes sense, with a regression test checking that the targets are preserved and predictions remain unchanged. This comment confirms the reproduction; the fix has not been pushed yet.

@sylvesterkaczmarek sylvesterkaczmarek left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rechecked current head 276ded8. This resolves the zero-column subsampling case I raised: subsampling is now skipped when there are no predictor columns in both implementations, and the regression covers positive max_rows_per_col values while confirming the padded targets are preserved and predictions remain zero. No remaining blocker from my review.

@rajpratham1 rajpratham1 left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The empty-dimension handling is correctly addressed in both XReg implementations, including safe padding for zero-sized dimensions and avoiding subsampling when there are zero predictor columns. The added regression tests cover empty shapes, normal padding, single-category fitting, and zero-predictor subsampling behavior. I don't see a blocking issue in this PR.

@rajpratham1 rajpratham1 left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The empty-dimension handling looks correct and is consistently applied to both src/ and v1/src/. The new tests cover zero-sized dimensions, existing non-empty padding behavior, and the zero-predictor subsampling case. I don't see any blocking issues in this PR.

@rajpratham1 rajpratham1 left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The empty-dimension handling is implemented consistently in both src and v1, and the added tests cover the affected zero-sized dimensions, normal padding behavior, and the zero-predictor subsampling path. I don't see any blocking issues in this PR.

@Bitshank-2338

Copy link
Copy Markdown
Author

Hi @rajatsen91, could you take a look at this fix for #197 when you have a chance?

The follow-up in 276ded8 addresses the zero-column subsampling issue raised during review in both XReg implementations. The regression tests cover positive max_rows_per_col values with debug_info=True, preserving the padded targets while keeping predictions unchanged. All 7 focused tests pass locally, and the CLA and check-changes checks pass. @sylvesterkaczmarek has also rechecked the updated commit and confirmed that their concern is resolved.

I'd appreciate a maintainer review and consideration for merge. Please let me know if any further changes are needed. Thank you!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

empty multidimensional numpy arrays causing error

3 participants