Simplify axis value checks

[zkneupper]

Simplify axis value checks

Example: replace axis != 1 and axis != 2 with axis not in [1, 2]

[mthrok]

Hi @zkneupper

Thanks for the submission. This is very common pattern in Python but this code is also executed by TorchScript runtime. Let us check the performance implication first.

@carolineechen temporarily assigning you. Let's talk about the benchmarking later.

[zkneupper]

Hi @zkneupper

Thanks for the submission. This is very common pattern in Python but this code is also executed by TorchScript runtime. Let us check the performance implication first.

@carolineechen temporarily assigning you. Let's talk about the benchmarking later.

Hi @zkneupper

Thanks for the submission. This is very common pattern in Python but this code is also executed by TorchScript runtime. Let us check the performance implication first.

@carolineechen temporarily assigning you. Let's talk about the benchmarking later.

Thanks @mthrok & @carolineechen. Are there tests within the repository for performance benchmarking?

[mthrok]

Thanks @mthrok & @carolineechen. Are there tests within the repository for performance benchmarking?

No, there is not. It's simply running the necessary function in timer. It's something like the following.

---

OMP_NUM_THREADS=1 numactl --membind 0 --cpubind 0 python -m timeit -n 1 -r 5 -s \
'import torch; a = torch.rand(1000, 1000)' \
'a.sum(1)'

This will set the number of threads used to 1 (OMP_NUM_THREADS=1), run the process on NUMA node 0 and memory node 0 (so the memory is being allocated right next to the CPU), runs a.sum(1) once (-n 1) within a benchmark and repeats (-r 5) the benchmark 5 times and executes import torch; a = torch.rand(1000, 1000) at startup.

---

For this PR, we need to benchmark the function in TorchScript, so in initialization, we first script the function, fun = torch.jit.script(torchaudio.functional.XXX) then run the operation func(...). If you are curious, you can try it before/after the change and give us the report too.

[mthrok]

Verified that there is no performance degradation in this change with TorchScript.