feat(decompression): add numpy and scipy output formats#142
Open
adrhill wants to merge 5 commits into
Open
Conversation
Add numpy_dense, scipy_coo, scipy_csr, and scipy_csc as output format options for Jacobian and Hessian decompression. - Extend OutputFormat type alias with new formats - Add scipy as optional dependency (pip install asdex[scipy]) - Add _to_scipy_sparse helper with lazy import and helpful error message - Add _convert_pytree_to_format for converting JAX arrays to target format - Update docstrings for all 8 public decompression functions - Add test fixtures: scipy_output_format, all_output_format - Add 11 explicit type-checking tests for new formats Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Owner
Author
|
CC @sirmarcel |
There was a problem hiding this comment.
Pull request overview
This PR extends asdex’s Jacobian/Hessian decompression APIs to support additional output formats backed by NumPy and SciPy, including adding SciPy as an optional dependency and ensuring missing-SciPy scenarios produce a helpful install message.
Changes:
- Adds new
output_formatoptions:numpy_dense,scipy_coo,scipy_csr,scipy_csc(and updates validation/error messaging). - Implements post-assembly conversion utilities in decompression to produce NumPy/SciPy outputs (with lazy SciPy import and a guided ImportError).
- Expands the test suite and pytest fixtures to cover the new output formats and SciPy-optional test behavior.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_decompression.py | Adds tests for NumPy/SciPy output formats and expected behavior. |
| tests/conftest.py | Extends dense-conversion helper and adds fixtures for SciPy/all output formats. |
| src/asdex/modes.py | Extends OutputFormat literal and updates output-format validation messaging/docs. |
| src/asdex/decompression.py | Adds NumPy/SciPy conversion helpers and hooks them into Jacobian/Hessian builds. |
| pyproject.toml | Adds SciPy as an optional extra (asdex[scipy]) and as a dev dependency. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #142 +/- ##
==========================================
- Coverage 93.14% 92.95% -0.20%
==========================================
Files 53 53
Lines 3518 3565 +47
==========================================
+ Hits 3277 3314 +37
- Misses 241 251 +10 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
- Use match statements in _convert_leaf_to_format - Remove pytest.importorskip from scipy fixtures (tests should fail loudly) - Parametrize scipy output format tests to reduce duplication - Use noqa comments for in-function scipy imports Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- Use match/assert_never in _to_scipy_sparse and _convert_leaf_to_format - Add Literal type aliases _JaxOutputFormat, _NumpyOutputFormat, _ScipyOutputFormat - Define OutputFormat as union of these types - Add ValueError for scipy formats with non-2D arrays - Consolidate output format tests using all_output_format fixture - Add numpy_dense to e2e test parametrization - Remove unused scipy_output_format fixture Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Add tests that verify SciPy formats raise ValueError for: - High-dimensional Hessian blocks (multi-dim inputs) - PyTree outputs (which produce non-2D Jacobian blocks) Also improve OutputFormat docstring to document the 2D constraint. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Update all jacobian/hessian function docstrings to note that SciPy formats only support 2D arrays (flat inputs and outputs). Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Comment on lines
1149
to
+1165
| @@ -985,7 +1158,12 @@ def _build_jacobian( | |||
|
|
|||
| # General path: scatter to dense, then assemble blocks. | |||
| dense = _scatter_dense(coloring, data) | |||
| return _assemble_jacobian(dense, sparsity, output_format, out_struct) | |||
| jac = _assemble_jacobian(dense, sparsity, output_format, out_struct) | |||
|
|
|||
| # Convert to numpy/scipy formats after assembly | |||
| if output_format in ("numpy_dense", "scipy_coo", "scipy_csr", "scipy_csc"): | |||
| return _convert_pytree_to_format(jac, output_format) | |||
Comment on lines
1181
to
+1194
| # Fast path: single input leaf with BCOO format and trivial pytree structure. | ||
| # numpy/scipy formats involve host transfer anyway, so fromdense cost is negligible. | ||
| if output_format == "bcoo" and _is_simple_input(sparsity): | ||
| in_shape = sparsity.leaf_shapes[0] | ||
| return sparsity.to_bcoo(data=data).reshape((*in_shape, *in_shape)) | ||
|
|
||
| # General path: scatter to dense, then assemble blocks. | ||
| dense = _scatter_dense(coloring, data) | ||
| return _assemble_hessian(dense, sparsity, output_format) | ||
| hess = _assemble_hessian(dense, sparsity, output_format) | ||
|
|
||
| # Convert to numpy/scipy formats after assembly | ||
| if output_format in ("numpy_dense", "scipy_coo", "scipy_csr", "scipy_csc"): | ||
| return _convert_pytree_to_format(hess, output_format) | ||
| return hess |
Comment on lines
1160
to
+1165
| dense = _scatter_dense(coloring, data) | ||
| return _assemble_jacobian(dense, sparsity, output_format, out_struct) | ||
| jac = _assemble_jacobian(dense, sparsity, output_format, out_struct) | ||
|
|
||
| # Convert to numpy/scipy formats after assembly | ||
| if output_format in ("numpy_dense", "scipy_coo", "scipy_csr", "scipy_csc"): | ||
| return _convert_pytree_to_format(jac, output_format) |
Comment on lines
+45
to
+51
| @pytest.fixture( | ||
| params=["dense", "bcoo", "numpy_dense", "scipy_coo", "scipy_csr", "scipy_csc"] | ||
| ) | ||
| def all_output_format(request): | ||
| """Parametrize over all output formats.""" | ||
| return request.param | ||
|
|
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
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.
Summary
numpy_dense,scipy_coo,scipy_csr, andscipy_cscas output format options for Jacobian and Hessian decompressionpip install asdex[scipy])New output formats
numpy_densenumpy.ndarrayscipy_cooscipy.sparse.coo_arrayscipy_csrscipy.sparse.csr_arrayscipy_cscscipy.sparse.csc_arrayTest plan
scipy_output_formatandall_output_formattest fixtures🤖 Generated with Claude Code