Optimize einsum_sparse using np.einsum_path

[apocalypse9949]

Problem

The einsum_sparse implementation in econml/utilities.py was merging sparse tensor dimensions using an arbitrary left-to-right reduction order. This could create unnecessarily large intermediate tensors, leading to excessive memory allocations and computations for complex sparse einsum operations. There was already a TODO indicating that the contraction order could be optimized.

Solution

• Updated einsum_sparse to utilize np.einsum_path.
• Created dummy arrays with shape information and zero sizes internally, ensuring np.einsum_path evaluates the optimal memory and time complexity for a given contraction without actually allocating heavy intermediate objects.
• Adapted the merging loop to respect the parsed optimal order output of np.einsum_path, rather than doing a fixed pop.
• Preserved the original API and fall back correctly to the left-to-right approach in error scenarios.

Validation

• Testing: Added test_einsum_sparse_optimization to verify correctness. It performs an optimized sparse operation against dense equivalents via np.einsum and validates the resulting dense outputs map equivalently.
• Executed relevant tests locally (pytest econml/tests/test_utilities.py and econml/tests/test_two_stage_least_squares.py), ensuring complete accuracy with no regressions.
• Benchmarking: Ran local performance checks locally (using a generic non-optimized 4-argument einsum string), yielding visible computational acceleration for larger dimensions.