Replace custom axis handling with numpy.array_utils.normalize_axis_*

[7astro7]

Addresses issue #2816

• Replaced manual axis checks with:
  • normalize_axis_index in partial_trace.py and partial_transpose.py
  • normalize_axis_index and normalize_axis_tuple in axis_atom.py
• Removed redundant ValueError checks for out-of-bounds axes
• Removed tests that manually validated axis bounds (since NumPy now handles this)

Description

Please include a short summary of the change. Issue link (if applicable): Please see above.

Type of change

• [x] New feature (backwards compatible)
• [ ] New feature (breaking API changes)
• [ ] Bug fix
• [ ] Other (Documentation, CI, ...)

Contribution checklist

• [x] Add our license to new files.
• [x] Check that your code adheres to our coding style.
• [x] Write unittests.
• [x] Run the unittests and check that they’re passing.
• [] Run the benchmarks to make sure your change doesn’t introduce a regression.

FYI I did not see the benchmark test.

[CLAassistant]

All committers have signed the CLA.

[7astro7]

@Transurgeon Thanks for the quick review. Would you prefer:

1. Bumping the numpy version in pyproject.toml to >=1.26 Release notes , or
2. Tabling this change for now?

Happy to update the dependencies if that direction works.

[Transurgeon]

@Transurgeon Thanks for the quick review. Would you prefer:

1. Bumping the numpy version in pyproject.toml to >=1.26 Release notes , or
2. Tabling this change for now?

Happy to update the dependencies if that direction works.

I think we can update the dependencies. But let's also hold off until after releasing 1.7.0 (which will be in a month or so?).

Thanks for your PR!

[7astro7]

@Transurgeon Sounds good, will revisit after 1.7.0 release.

[7astro7]

@Transurgeon Rebased onto latest master and bumped the NumPy runtime requirement to >=1.26.0 to support normalize_axis_index. Let me know if you’d like any additional adjustments.

[7astro7]

@Transurgeon Would you like to close this to clear out noise or keep it open for further discussion?

[PTNobel]

NumPy 1.25 is EoL so I'm fine with the 1.26 dependency.

https://endoflife.date/numpy

[github-actions[bot]]

Benchmarks that have improved:

   before           after         ratio
 [9b1f4628]       [cfce3d6e]

•     1.14±0s          917±0ms     0.80  gini_portfolio.Cajas.time_compile_problem
  

Benchmarks that have stayed the same:

   before           after         ratio
 [9b1f4628]       [cfce3d6e]
      845±0ms          874±0ms     1.03  simple_QP_benchmarks.LeastSquares.time_compile_problem
      1.11±0s          1.12±0s     1.02  finance.FactorCovarianceModel.time_compile_problem
      299±0ms          303±0ms     1.01  slow_pruning_1668_benchmark.SlowPruningBenchmark.time_compile_problem
      345±0ms          350±0ms     1.01  gini_portfolio.Yitzhaki.time_compile_problem
      22.4±0s          22.7±0s     1.01  sdp_segfault_1132_benchmark.SDPSegfault1132Benchmark.time_compile_problem
      4.57±0s          4.60±0s     1.01  huber_regression.HuberRegression.time_compile_problem
      1.86±0s          1.87±0s     1.01  simple_QP_benchmarks.UnconstrainedQP.time_compile_problem
      723±0ms          727±0ms     1.01  matrix_stuffing.ConeMatrixStuffingBench.time_compile_problem
      1.37±0s          1.37±0s     1.00  matrix_stuffing.ParamConeMatrixStuffing.time_compile_problem
      235±0ms          236±0ms     1.00  gini_portfolio.Murray.time_compile_problem
      2.34±0s          2.35±0s     1.00  simple_LP_benchmarks.SimpleFullyParametrizedLPBenchmark.time_compile_problem
      285±0ms          285±0ms     1.00  matrix_stuffing.ParamSmallMatrixStuffing.time_compile_problem
      5.04±0s          5.00±0s     0.99  optimal_advertising.OptimalAdvertising.time_compile_problem
      255±0ms          253±0ms     0.99  simple_QP_benchmarks.SimpleQPBenchmark.time_compile_problem
      5.20±0s          5.15±0s     0.99  svm_l1_regularization.SVMWithL1Regularization.time_compile_problem
      1.69±0s          1.67±0s     0.99  tv_inpainting.TvInpainting.time_compile_problem
      566±0ms          560±0ms     0.99  simple_QP_benchmarks.ParametrizedQPBenchmark.time_compile_problem
      5.35±0s          5.28±0s     0.99  semidefinite_programming.SemidefiniteProgramming.time_compile_problem
      11.5±0s          11.4±0s     0.99  simple_LP_benchmarks.SimpleLPBenchmark.time_compile_problem
      243±0ms          239±0ms     0.98  high_dim_convex_plasticity.ConvexPlasticity.time_compile_problem
      13.6±0s          13.3±0s     0.98  finance.CVaRBenchmark.time_compile_problem
      961±0ms          940±0ms     0.98  simple_LP_benchmarks.SimpleScalarParametrizedLPBenchmark.time_compile_problem
      2.96±0s          2.89±0s     0.98  quantum_hilbert_matrix.QuantumHilbertMatrix.time_compile_problem
     46.3±0ms         44.5±0ms     0.96  matrix_stuffing.SmallMatrixStuffing.time_compile_problem

[PTNobel]

Looks like we'd need NumPy 2.0 to get the desired functions tho? Looking at the CI

[7astro7]

To keep compatibility with both 1.26 and 2.x a small compatibility shim could solve this.

try:
    from numpy.lib.array_utils import normalize_axis_index, normalize_axis_tuple  # numpy ≥ 2.0
except ModuleNotFoundError:
    from numpy.core.multiarray import normalize_axis_index, normalize_axis_tuple  # numpy < 2.0

That way the import works for numpy >= 1.26. Let me know if you’d like me to push this change to the PR

[PTNobel]

Before we do the shim, let me see if we have consensus to move to NumPy 2.0.

@SteveDiamond @Transurgeon shall we merge this PR bumping our dependencies to NumPy 2.0 this month, in compliance with Spec 0? We can also bump to Python 3.11 in a different PR and add some docs on our dependencies pointing to Spec 0.

[Transurgeon]

Before we do the shim, let me see if we have consensus to move to NumPy 2.0.

@SteveDiamond @Transurgeon shall we merge this PR bumping our dependencies to NumPy 2.0 this month, in compliance with Spec 0? We can also bump to Python 3.11 in a different PR and add some docs on our dependencies pointing to Spec 0.

I think we should! Let's make sure to mark these bumps as "PR no backport needed" so that 1.7.x can still provide wheels for py3.10 and support for numpy < 2.0. What do you think?

[PTNobel]

I think we should! Let's make sure to mark these bumps as "PR no backport needed" so that 1.7.x can still provide wheels for py3.10 and support for numpy < 2.0. What do you think?

Sounds good to me.

[Transurgeon]

Okay I think we should merge this! Thanks again for your help and patience @7astro7 !