Adding vmap support for special bessel functions

[SiddharthIVEX]

Fixes #91402

Description

This PR adds vmap support for the following bessel functions under torch.special.

• special.bessel_j0
• special.bessel_y0
• special.bessel_j1
• special.modified_bessel_i0
• special.bessel_y1
• special.scaled_modified_bessel_k0
• special.scaled_modified_bessel_k1
• special.modified_bessel_i1

Files changed:

1. aten/src/ATen/functorch/BatchRulesUnaryOps.cpp
2. test/functorch/test_vmap.py

How was the PR tested?

1. The unit tests under test_vmap.py were run and all of them passed. The output is shown below.

configfile: pytest.ini
plugins: hypothesis-6.71.0, anyio-2.2.0
collected 2003 items / 1981 deselected / 22 selected                                                                                                                 

test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_special_bessel_j0_cpu_float32 PASSED                                           [  4%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_special_bessel_j1_cpu_float32 PASSED                                           [  9%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_special_bessel_y0_cpu_float32 PASSED                                           [ 13%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_special_bessel_y1_cpu_float32 PASSED                                           [ 18%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_special_modified_bessel_i0_cpu_float32 PASSED                                  [ 22%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_special_modified_bessel_i1_cpu_float32 PASSED                                  [ 27%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_special_modified_bessel_k0_cpu_float32 PASSED                                  [ 31%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_special_modified_bessel_k1_cpu_float32 PASSED                                  [ 36%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_special_scaled_modified_bessel_k0_cpu_float32 PASSED                           [ 40%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_special_scaled_modified_bessel_k1_cpu_float32 PASSED                           [ 45%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_special_spherical_bessel_j0_cpu_float32 PASSED                                 [ 50%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_vmap_exhaustive_special_bessel_j0_cpu_float32 PASSED                                             [ 54%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_vmap_exhaustive_special_bessel_j1_cpu_float32 PASSED                                             [ 59%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_vmap_exhaustive_special_bessel_y0_cpu_float32 PASSED                                             [ 63%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_vmap_exhaustive_special_bessel_y1_cpu_float32 PASSED                                             [ 68%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_vmap_exhaustive_special_modified_bessel_i0_cpu_float32 PASSED                                    [ 72%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_vmap_exhaustive_special_modified_bessel_i1_cpu_float32 PASSED                                    [ 77%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_vmap_exhaustive_special_modified_bessel_k0_cpu_float32 PASSED                                    [ 81%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_vmap_exhaustive_special_modified_bessel_k1_cpu_float32 PASSED                                    [ 86%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_vmap_exhaustive_special_scaled_modified_bessel_k0_cpu_float32 PASSED                             [ 90%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_vmap_exhaustive_special_scaled_modified_bessel_k1_cpu_float32 PASSED                             [ 95%]
test/functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_vmap_exhaustive_special_spherical_bessel_j0_cpu_float32 PASSED                                   [100%]

================================================================ 22 passed, 1981 deselected in 18.42s ================================================================

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ROCm jobs are failing with No space left on device

✅ No Failures

As of commit 93f85c8:
💚 Looks good so far! There are no failures yet. 💚

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

[zou3519]

LGTM, thank you!

[zou3519]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here