Dispatch numpy.take_along_axis to torch.take_along_dim

[guilhermeleobas]

Stack from ghstack (oldest at bottom):

• Remove include_0d from sample_inputs_gather #109125
• Update take_along_dim docs to include dim=None case #109120
• -> Dispatch numpy.take_along_axis to torch.take_along_dim #108880
• Add SymInt support to torch.take_along_dim #108879

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng

[pytorch-bot]

🔗 Helpful Links

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

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

✅ You can merge normally! (2 Unrelated Failures)

As of commit 1173859 with merge base bde75eb ():

UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:

• linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 3, 5, linux.g5.4xlarge.nvidia.gpu, unstable) (gh)
• linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 4, 5, linux.g5.4xlarge.nvidia.gpu, unstable) (gh)

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

[lezcano]

What's the issue with the empty one?

[guilhermeleobas]

It fails in a test case on test_misc.py as the indices has dtype uint8_t:

...
torch/_dynamo/utils.py:1339: in <lambda>
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)
torch/_dynamo/utils.py:1411: in run_node
    raise RuntimeError(fn_str + str(e)).with_traceback(e.__traceback__) from e
torch/_dynamo/utils.py:1398: in run_node
    return node.target(*args, **kwargs)
torch/_dynamo/utils.py:1699: in __call__
    out = self.f(*args, **kwargs)
torch/_numpy/_normalizations.py:213: in wrapped
    result = func(*args, **kwds)
torch/_numpy/_funcs_impl.py:889: in take_along_axis
    return torch.take_along_dim(arr, indices, axis)
E   torch._dynamo.exc.TorchRuntimeError: Failed running call_function <Wrapped function <original take_along_axis>>(*(FakeTensor(..., size=(s0,)), FakeTensor(..., size=(0,), dtype=torch.uint8), 0), **{}):
E   torch.take_along_dim(): dtype of indices should be Long but got Byte
E
E   from user code:
E      File "/home/guilhermeleobas/git/pytorch/test/dynamo/test_misc.py", line 1283, in fn
E       return np.take_along_axis(x, i, a)
E
E   Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
E
E
E   You can suppress this exception and fall back to eager by setting:
E       import torch._dynamo
E       torch._dynamo.config.suppress_errors = True
E
E
E   To execute this test, run the following from the base repo dir:
E        python test/dynamo/test_misc.py -k test_numpy_take_along_axis
E
E   This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0

[lezcano]

Makes sense. Actually, it may be best to use sample_inputs_take_along_dim. Note that the order of the args for indices and dims is swapped wrt. gather.

[guilhermeleobas]

Are you ok if I skip the tests from sample_inputs_take_along_dim that don't include dim?

[lezcano]

What do you mean exactly?

[guilhermeleobas]

Sorry for not including more information in my previous comment. There are two tests that don't include the dim argument in sample_inputs_take_along_dim

https://github.com/pytorch/pytorch/blob/f9a250c35bd061e2e6f4c2d92e2b1b16390e8636/torch/testing/_internal/common_methods_invocations.py#L2842-L2846

[lezcano]

Actually, Numpy also has this behaviour, so we should also implement it. Can you modify the decomposition accordingly?

Also, could you send a separate PR fixing the docs of this operation noting that dim can also be None?

[lezcano]

Just a small note, feel free to merge once you change the sampling function.
You can even clean-up the include_0d and the include_empty args for sample_inputs_gather as these will not be used anywhere any more.

Feel free to merge once you change this. In a stack, if you call pytorchbot merge on a given PR, it'll merge that PR and all the previous ones.

[lezcano]

Makes sense. Actually, it may be best to use sample_inputs_take_along_dim. Note that the order of the args for indices and dims is swapped wrt. gather.

[lezcano]

Also, note that you can even clean-up the include_0d and the include_empty args for sample_inputs_gather as these will not be used anywhere any more.