Improve tabulate

[cgarciae]

What does this PR do?

Overhauls tabulate with a new system to capture call information and fixes #2274 and #2359.

Changes

• Adds a new thread-local _CallInfoContext context manager that provides richer call information.
• Table entries are now sorted by call order which is more intuitive / correct.
• Two new columns are added to the table:
  • module: shows module type
  • inputs: shows input information
• Removes the exclude_methods arguments, it is no longer necessary as capture_intermediates is no longer used.
  • This fixes nn.tabulate results in KeyError: 'intermediates' with methods that include transformations #2274 as collecting summaries no longer depends on any variables collection.
• A new show_repeated arguments is added, lets you decide if you want to show repeated modules in case the same module is called multiple times.
• Adds a new console_kwargs argument that lets users configure the rich.console.Console object used to render the table.

Future issues

Some issues related to Scope.path where uncovered here, will post issues with a minimal repro in the near future. To get around these issues some temporal fixes where used, see TODOs associated with each issue.

• Scope.path sometimes empty when a module is called more than once, see TODO.
• Scope.path names sometimes don't match the variable structure, specific issue found when naming a lifted Module, see TODO.

Sample

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 97.83784% with 4 lines in your changes missing coverage. Please review.

Project coverage is 70.00%. Comparing base (f754541) to head (b1025cc).
Report is 1757 commits behind head on main.

Files with missing lines	Patch %	Lines
flax/linen/module.py	96.22%	2 Missing ⚠️
flax/linen/summary.py	98.48%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2316      +/-   ##
==========================================
+ Coverage   69.33%   70.00%   +0.67%     
==========================================
  Files          63       63              
  Lines        5556     5664     +108     
==========================================
+ Hits         3852     3965     +113     
+ Misses       1704     1699       -5     

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

[jheek]

Looking good!

[marcvanzee]

@jheek Something seems to have gone wrong with this PR internally, could you please take a look?