experiment: support for uma#319
Conversation
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThis PR extends Pyodide Torch patching to support FAIRChem/UMA, adds a ChangesFAIRChem UMA Support in Pyodide
🎯 4 (Complex) | ⏱️ ~45 minutes
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@other/experiments/jupyterlite/relax_structure_with_uma.ipynb`:
- Around line 64-80: The notebook cell is malformed: remove the duplicated
"cell_type": "code" key so the JSON object has a single cell_type entry, and fix
the Python indentation by trimming the leading spaces on the lines importing and
calling apply_all_patches (the lines referencing
mat3ra.notebooks_utils.pyodide.packages.torch and apply_all_patches) so they are
left-aligned with the other statements; ensure the cell's "source" array
contains only correctly-indented strings and valid JSON keys.
In `@src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py`:
- Around line 569-571: The file
src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py was reformatted by
Black; run the formatter (e.g., run `black
src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py` or `pre-commit run
--all-files`) and commit the updated file so the assignment to
sys.modules["ray.util.scheduling_strategies"].PlacementGroupSchedulingStrategy =
_PlacementGroupSchedulingStrategy matches Black's formatting; then push the
formatted changes.
- Around line 517-518: The lambda for numba_mod.njit currently returns the same
decorator in both branches, breaking direct usage like `@njit`; change the
conditional so that when a first positional argument is present and is callable
(i.e., the decorator used directly as `@njit`), `njit` returns that function
immediately, otherwise it returns a decorator function; update the expression
for numba_mod.njit accordingly and keep numba_mod.jit = numba_mod.njit to
preserve aliasing.
🪄 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: 1bded097-58b6-42c4-9e82-3995f373b0e1
📒 Files selected for processing (3)
config.ymlother/experiments/jupyterlite/relax_structure_with_uma.ipynbsrc/py/mat3ra/notebooks_utils/pyodide/packages/torch.py
| numba_mod.njit = lambda *a, **k: (lambda f: f) if not a or callable(a[0]) else lambda f: f | ||
| numba_mod.jit = numba_mod.njit |
There was a problem hiding this comment.
Bug: Both ternary branches return identical values, breaking @njit direct decoration.
The condition (lambda f: f) if not a or callable(a[0]) else lambda f: f returns lambda f: f in both branches. When used as @njit (without parentheses), njit(fn) is called and should return fn directly, not a decorator.
🐛 Proposed fix
- numba_mod.njit = lambda *a, **k: (lambda f: f) if not a or callable(a[0]) else lambda f: f
+ numba_mod.njit = lambda *a, **k: a[0] if a and callable(a[0]) else lambda f: f📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| numba_mod.njit = lambda *a, **k: (lambda f: f) if not a or callable(a[0]) else lambda f: f | |
| numba_mod.jit = numba_mod.njit | |
| numba_mod.njit = lambda *a, **k: a[0] if a and callable(a[0]) else lambda f: f | |
| numba_mod.jit = numba_mod.njit |
🧰 Tools
🪛 Ruff (0.15.13)
[warning] 517-517: Useless if-else condition
(RUF034)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py` around lines 517 -
518, The lambda for numba_mod.njit currently returns the same decorator in both
branches, breaking direct usage like `@njit`; change the conditional so that
when a first positional argument is present and is callable (i.e., the decorator
used directly as `@njit`), `njit` returns that function immediately, otherwise
it returns a decorator function; update the expression for numba_mod.njit
accordingly and keep numba_mod.jit = numba_mod.njit to preserve aliasing.
There was a problem hiding this comment.
♻️ Duplicate comments (1)
src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py (1)
583-585:⚠️ Potential issue | 🟠 Major | ⚡ Quick winFix
numba_mod.njitdirect-decorator behavior.This stub returns a decorator in both branches; direct usage (
@njit/njit(fn)) should returnfnimmediately, otherwise decorated call paths can break at runtime.🐛 Proposed fix
- numba_mod.njit = lambda *a, **k: (lambda f: f) if not a or callable(a[0]) else lambda f: f + numba_mod.njit = lambda *a, **k: a[0] if a and callable(a[0]) else (lambda f: f)#!/bin/bash # Verify the faulty expression is still present. rg -nP 'numba_mod\.njit\s*=\s*lambda \*a, \*\*k: \(lambda f: f\) if not a or callable\(a\[0\]\) else lambda f: f'🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py` around lines 583 - 585, The stub numba_mod.njit currently always returns a decorator, breaking direct-usage patterns like `@njit` or njit(fn); change njit to return the callable immediately when invoked with a single callable positional arg and no kwargs, otherwise return the identity decorator. Concretely, replace the lambda with a small function (named njit or assigned to numba_mod.njit) that does: if len(a) == 1 and callable(a[0]) and not k: return a[0] else: return (lambda f: f); keep numba_mod.jit assigned to the same implementation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Duplicate comments:
In `@src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py`:
- Around line 583-585: The stub numba_mod.njit currently always returns a
decorator, breaking direct-usage patterns like `@njit` or njit(fn); change njit to
return the callable immediately when invoked with a single callable positional
arg and no kwargs, otherwise return the identity decorator. Concretely, replace
the lambda with a small function (named njit or assigned to numba_mod.njit) that
does: if len(a) == 1 and callable(a[0]) and not k: return a[0] else: return
(lambda f: f); keep numba_mod.jit assigned to the same implementation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4936510c-3730-43c3-b257-acde9062f8c6
📒 Files selected for processing (4)
config.ymlother/experiments/jupyterlite/relax_structure_with_uma.ipynbpackages/models/uma-s-1p1-int8.ptsrc/py/mat3ra/notebooks_utils/pyodide/packages/torch.py
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py (1)
823-851: 💤 Low valueConsider guarding the
fairchemimport in INT8 dequantization path.If
torch.loadis called on an INT8 checkpoint before FAIRChem is installed, the import on line 827 will raiseImportError. This monkey-patch runs unconditionally whenpatch_fairchem_deps()is called.Potential guard
if isinstance(result, dict) and "quantized_ema_state_dict" in result: + try: + from fairchem.core.units.mlip_unit.api.inference import MLIPInferenceCheckpoint + except ImportError: + print("⚠ INT8 checkpoint detected but fairchem-core not installed; returning raw dict") + return result import gc as _gc - from fairchem.core.units.mlip_unit.api.inference import MLIPInferenceCheckpoint🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py` around lines 823 - 851, The monkey-patched loader _int8_aware_torch_load unconditionally imports MLIPInferenceCheckpoint from fairchem which raises ImportError if fairchem isn't installed; wrap the fairchem import (the line importing MLIPInferenceCheckpoint inside _int8_aware_torch_load) in a try/except ImportError and handle the missing dependency by skipping dequantization and returning the original result (or otherwise returning a sensible fallback) instead of raising—ensure torch.load remains safe when patch_fairchem_deps() runs but fairchem is absent.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py`:
- Around line 823-851: The monkey-patched loader _int8_aware_torch_load
unconditionally imports MLIPInferenceCheckpoint from fairchem which raises
ImportError if fairchem isn't installed; wrap the fairchem import (the line
importing MLIPInferenceCheckpoint inside _int8_aware_torch_load) in a try/except
ImportError and handle the missing dependency by skipping dequantization and
returning the original result (or otherwise returning a sensible fallback)
instead of raising—ensure torch.load remains safe when patch_fairchem_deps()
runs but fairchem is absent.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4e2ca0a5-5bf6-48a1-9000-37dba3a6e0f2
📒 Files selected for processing (4)
.gitignoreconfig.ymlother/experiments/jupyterlite/relax_structure_with_uma.ipynbsrc/py/mat3ra/notebooks_utils/pyodide/packages/torch.py
✅ Files skipped from review due to trivial changes (1)
- .gitignore
🚧 Files skipped from review as they are similar to previous changes (1)
- config.yml
1 participant
Summary by CodeRabbit
New Features
New Features / Improvements
Chores