Skip to content

Use grad to compute jacobian when input shape is known to be (1,) #1454

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 8, 2025

Conversation

jessegrabowski
Copy link
Member

@jessegrabowski jessegrabowski commented Jun 7, 2025

Description

Small optimization for pt.jacobian(x, *args). If x is known to be shape (1,), we should still use pt.grad rather than doing a scan.

Related Issue

  • Closes #
  • Related to #

Checklist

Type of change

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

📚 Documentation preview 📚: https://pytensor--1454.org.readthedocs.build/en/1454/

@jessegrabowski jessegrabowski added bug Something isn't working gradients labels Jun 7, 2025
@jessegrabowski jessegrabowski requested a review from ricardoV94 June 7, 2025 14:31
@ricardoV94
Copy link
Member

ricardoV94 commented Jun 7, 2025

PR title is imprecise, any array with statically known size 1 now uses this, so shape=(1, ..., 1)

Copy link

codecov bot commented Jun 7, 2025

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.12%. Comparing base (ff98ab8) to head (818b641).
Report is 1 commits behind head on main.

Additional details and impacted files

Impacted file tree graph

@@           Coverage Diff           @@
##             main    #1454   +/-   ##
=======================================
  Coverage   82.12%   82.12%           
=======================================
  Files         211      211           
  Lines       49757    49757           
  Branches     8819     8819           
=======================================
  Hits        40862    40862           
  Misses       6715     6715           
  Partials     2180     2180           
Files with missing lines Coverage Δ
pytensor/gradient.py 78.55% <100.00%> (ø)
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@jessegrabowski
Copy link
Member Author

PR title is imprecise, any array with statically known size 1 now uses this, so shape=(1, ..., 1)

No, because jacobian explicitly checks for ndim=1. So shape(1,1,1) will raise before we get to the change in this PR.

@jessegrabowski jessegrabowski merged commit 0ea61bc into pymc-devs:main Jun 8, 2025
74 checks passed
@jessegrabowski jessegrabowski deleted the jacobian-shape-check branch June 8, 2025 15:06
@ricardoV94
Copy link
Member

PR title is imprecise, any array with statically known size 1 now uses this, so shape=(1, ..., 1)

No, because jacobian explicitly checks for ndim=1. So shape(1,1,1) will raise before we get to the change in this PR.

I think I'm my PR on vectorized I removed that constraint?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working gradients
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants