Skip to content

Conversation

@mj023
Copy link
Collaborator

@mj023 mj023 commented Mar 13, 2025

There is an Error in the Jax QR Factorization that leads to slow execution for batches of small matrices when calculating the gradient. This PR tries to fix this, by implementing our own QR Factorization that runs much faster when using a GPU.

  • Implement faster QR
  • Identify the Jax backend and select the right QR Factorization

@codecov
Copy link

codecov bot commented Mar 13, 2025

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 91.68%. Comparing base (63b6ceb) to head (20976e5).
Report is 2 commits behind head on endogenize_investments.

Additional details and impacted files
@@                    Coverage Diff                     @@
##           endogenize_investments      #81      +/-   ##
==========================================================
+ Coverage                   91.49%   91.68%   +0.18%     
==========================================================
  Files                          40       42       +2     
  Lines                        3387     3464      +77     
==========================================================
+ Hits                         3099     3176      +77     
  Misses                        288      288              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@mj023 mj023 requested a review from hmgaudecker March 17, 2025 12:17
Copy link
Member

@hmgaudecker hmgaudecker left a comment

Choose a reason for hiding this comment

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

Looks great, thanks!

Haven't looked at the algorithm yet.

In addition to the couple of suggestions on the tests, can you please:

  • Set endogenize_investments as the target branch.
  • Run the pre-commit hooks on your changes (pixi global install pre-commit, pre-commit run --files src/skillmodels/kalman_filters.py etc.)

@mj023 mj023 changed the base branch from main to endogenize_investments March 17, 2025 15:20
@hmgaudecker hmgaudecker merged commit 9bc6d59 into endogenize_investments Mar 17, 2025
3 checks passed
@hmgaudecker hmgaudecker deleted the Fix_QR branch March 17, 2025 17:13
@hmgaudecker
Copy link
Member

Thanks a lot!

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.

3 participants