Remove redundant flatten/unflatten of PyTrees

[rauletorresc]

Context: When calling get_abstract_signature from within get_decomposed_signature we are flattening/unflattening the PyTree one more time without need. In this sense get_abstract_signature is just a wrapper for the shaped_abstractify function that we don't need here.

Description of the Change: Avoid the wrapper function call and use shaped_abstractify directly.

Benefits: Avoid code that does nothing.

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md on your branch with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.88%. Comparing base (1a6df92) to head (8e43f76).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1648   +/-   ##
=======================================
  Coverage   96.88%   96.88%           
=======================================
  Files          80       80           
  Lines        8861     8861           
  Branches      841      841           
=======================================
  Hits         8585     8585           
  Misses        222      222           
  Partials       54       54           

☔ 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.

[dime10]

This change looks good, but if we want to resolve the core issue we have to tackle it here:

catalyst/frontend/catalyst/jit.py

Line 715 in 3000919

dynamic_sig = get_abstract_signature(dynamic_args)

• We generate an abstract signature re-wrapped into pytrees.
• The question is whether we really need the abstract values at some point, and if so whether we need them in pytrees. We can modify either of those conditions to become compliant.

The abstract signature is used for example to generate jaxpr:

catalyst/frontend/catalyst/jit.py

Line 740 in 3000919

jaxpr, out_type, treedef, plugins = trace_to_jaxpr(

• but further down we already abstractify:

  catalyst/frontend/catalyst/jax_extras/tracing.py

  Line 530 in 3000919

  in_type, in_tree = abstractify(args, kwargs)

Then it is used again during promotion:

catalyst/frontend/catalyst/jit.py

Line 604 in 3000919

args = promote_arguments(self.c_sig, dynamic_args)

• Here we can probably just do the promotion on a flattened argument list instead.

The last use case, in the compilation cache, is already addressed in this PR.