Implemented Eye Op in PyTorch

[twaclaw]

Description

• Implemented Eye Op with diagonal offset in PyTorch
  - torch.eye doesn't support a diagonal offset param k directly. The implementation of this offset required composing two torch functions.
• Added tests for different dimensions and diagonal offset, including corner cases.

Related Issue

• Related to Implement all Ops in PyTorch (help welcome!) #821

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[codecov]

Codecov Report

Attention: Patch coverage is 76.92308% with 3 lines in your changes missing coverage. Please review.

Project coverage is 81.05%. Comparing base (ca10298) to head (08d23ab).
Report is 120 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/pytorch/dispatch/basic.py	76.92%	2 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #877      +/-   ##
==========================================
- Coverage   81.05%   81.05%   -0.01%     
==========================================
  Files         170      170              
  Lines       46980    46992      +12     
  Branches    11509    11510       +1     
==========================================
+ Hits        38079    38088       +9     
- Misses       6695     6697       +2     
- Partials     2206     2207       +1     

Files with missing lines	Coverage Δ
pytensor/link/pytorch/dispatch/basic.py	83.82% <76.92%> (-1.90%)	⬇️

[twaclaw]

@ricardoV94, is there something odd with the coverage reports?

[ricardoV94]

@ricardoV94, is there something odd with the coverage reports?

It's flaky don't worry

[ricardoV94]

In this case let's compile a single pytensor function and call repeatedly with the different inputs, instead of creating a function for every combination of test values. Something like:

fn = pytensor.function([N, M, k], out, mode=torch_mode)  # Whatever mode is used internally in `compare_pytorch_and_py`
for _N in trange....:
  ...
  np.testing.assert_allclose(fn(_N, _M, _k), np.eye(_N, _M, _k))

[jessegrabowski]

Do we need to specify dtype here? It should be handled automatically by the config.floatX setting

[twaclaw]

@ricardoV94 ?

[ricardoV94]

If dtype is something the user can control directly we should probably test with default and a user-specified non-default?

[jessegrabowski]

Left some comments, some are style choices (feel free to push back on those) but some need to be addressed (pytorch_mode), so I marked it as changes needed. The branch also needs to be rebased.

[twaclaw]

Can you give me a declaration example for the default, non-explicit part. Am Sa., 6. Juli 2024 um 19:05 Uhr schrieb Ricardo Vieira < ***@***.***>:

…

***@***.**** commented on this pull request. ------------------------------ In tests/link/pytorch/test_basic.py <#877 (comment)>: > @@ -235,3 +235,20 @@ def test_arange(): FunctionGraph([start, stop, step], [out]), [np.array(1), np.array(10), np.array(2)], ) + + +def test_eye(): + N = scalar("N", dtype="int64") + M = scalar("M", dtype="int64") + k = scalar("k", dtype="int64") + + out = eye(N, M, k, dtype="float32") + + trange = range(1, 6) + + fn = function([N, M, k], out, mode=pytorch_mode) + + for _N in trange: Can stay then, there's still the question of testing default vs manual dtype discussed above — Reply to this email directly, view it on GitHub <#877 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ALEEX7GXWIXDLW4V2THR7U3ZLAPVVAVCNFSM6AAAAABKIGQSLGVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDCNRRGU2TINBWGU> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[ricardoV94]

I guess just test dtype=config.floatX vs dtype="int64". Because that requires a different function I would use parametrize for that

[jessegrabowski]

Not sure why codecov is whining, but this looks good to go on my end.

[ricardoV94]

I think codecov cannot see line hits the way pytorch uses the functions

[ricardoV94]

Cool stuff @twaclaw!

[twiecki]

@twaclaw I've been trying to contact you through LI, or maybe send me an email.