Batching rules matrix exp

[Xiao215]

Fixes #115992

This is a follow up on PR #116377

Added matrix_exp to BatchRulesDecomposition.

cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/117934

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 15 New Failures, 1 Unrelated Failure

As of commit a2d1351 with merge base f1aef2c ():

NEW FAILURES - The following jobs have failed:

• Check mergeability of ghstack PR / ghstack-mergeability-check (gh)
  Process completed with exit code 1.
• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 2, 5, linux.4xlarge.nvidia.gpu) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• pull / linux-focal-py3.11-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• pull / linux-focal-py3.11-clang10 / test (default, 2, 3, linux.2xlarge) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• pull / linux-focal-py3.11-clang10 / test (dynamo, 2, 3, linux.2xlarge) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• pull / linux-focal-py3.8-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• pull / linux-focal-py3.8-clang10 / test (default, 2, 3, linux.2xlarge) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• pull / linux-focal-py3.8-clang10 / test (dynamo, 2, 3, linux.2xlarge) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• pull / linux-jammy-py3.10-clang15-asan / test (default, 4, 6, linux.4xlarge) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• pull / linux-jammy-py3.8-gcc11 / test (default, 2, 3, linux.2xlarge) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• trunk / linux-focal-cuda12.1-py3.10-gcc9 / test (nogpu_AVX512, 1, 1, linux.2xlarge) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• trunk / linux-focal-cuda12.1-py3.10-gcc9 / test (nogpu_NO_AVX2, 1, 1, linux.2xlarge) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• trunk / macos-12-py3-arm64 / test (default, 2, 3, macos-m1-12) (gh)
  functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]
• trunk / win-vs2019-cpu-py3 / test (default, 1, 3, windows.4xlarge.nonephemeral) (gh)
  functorch\test_vmap_registrations.py::TestFunctorchDispatcher::test_register_a_batching_rule_for_composite_implicit_autograd_[aten::matrix_exp]

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / linux-focal-cuda12.1-py3.10-gcc9-bazel-test / build-and-test (default, 1, 1, linux.4xlarge.nvidia.gpu) (gh)
  Build completed, 1 test FAILED, 46 total actions

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Xiao215]

Anyone knows why I got an Error: Process completed with exit code 128.?

[zou3519]

Anyone knows why I got an Error: Process completed with exit code 128.?

That one is unrelated, we can ignore it.

[zou3519]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Successfully rebased batching-rules-matrix_exp onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout batching-rules-matrix_exp && git pull --rebase)

[zou3519]

Please shout in three days if I haven't merged this yet

[zou3519]

test failures look real

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

[ZelboK]

@zou3519

Curious to know why you think this is failling? Locally I get this,

File "/home/ksm/pytorch/test/functorch/test_vmap_registrations.py", line 308, in test_unimplemented_batched_registrations
    assert registration in FuncTorchBatchedDecompositionRegistrations, (
AssertionError: Please check that there is an OpInfo that covers the operator aten::matrix_exp and add a registration in BatchedRulesDecompositions.cpp. If your operator isn't user facing, please add it to the xfail list

which is different from what the pipeline describes.