This is ready to review.

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 87.26%. Comparing base (78a45f6) to head (15be459).

@@            Coverage Diff             @@
##             main     #564      +/-   ##
==========================================
+ Coverage   87.25%   87.26%   +0.01%     
==========================================
  Files         125      125              
  Lines        7916     7908       -8     
==========================================
- Hits         6907     6901       -6     
+ Misses       1009     1007       -2     

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

Great job @aarmey! I haven't done a super thorough review, but it looks great at first glance!

I have one question. I see you're using Union- and Optional-types instead of the new |- and | None-syntax. I assume this is to support Python 3.9, which is EOL next year?

A workaround if we want to use the new syntax is adding a from __future__ import annotations, which turns all type annotations into strings. The upside is that it makes type annotations forward compatible, and a downside is that arbitrary syntax is allowed so typos might not be found. However IDEs are usually good at spotting typos, so I think it's worth it for the increased readibility.

Also, the previous default tolerance of 1e-8 was selected based on @MarieRoald's experience with PARAFAC2, so if we change the tolerance function, it might be worthwhile to also check that the default value is sensible (and works for non-simulated data also), but you might have done that already?

Thanks @aarmey sorry just saw you marked it as ready for review.

It looks good, just left some comments - I somehow find type hints more difficult to parse than just signature + doc, especially when it gets nested like Optional[Union[..]...] but I guess it can be nice to have for type check.

I agree that the type annotations can become unwieldy, @JeanKossaifi. I have found them to still be helpful, on balance, because they describe the argument types in a very exact manner:

1. Occasionally, the documentation is a little vague because arguments such as nn_mode are quite complicated. A more advanced user can inspect these to know exactly what can be passed in.
2. These seem to help with autocomplete tools, probably because they are machine-readable.
3. We have a complicated codebase, and I've found automated checkers such as ruff/pyright to be helpful. For instance, quickly skimming the results, tl.linalg.higher_order_moment() currently does not work because batched_outer() only takes one argument, not two. It's easy to overdo it with automated tools, but I wonder if these plus some type annotation could help to avoid a separate set of errors compared to our tests. (I have little experience with large codebases, so I am genuinely wondering what people think about this.)

These are all good points @aarmey, and I guess type annotations are getting better (especially in 1.12+) so things will get less cumbersome. Should we remove Python 3.9 and move to the newer type as @yngvem was suggesting (we can have a clear warning in the doc)?

Thanks, @JeanKossaifi and @yngvem! I like this plan, so I'll merge this and then follow up with a PR removing Python 3.9 support and simplifying the type annotations.