Expand distributed indexing, match numpy indexing scheme #316

	name: Benchmark PR
	on:
	pull_request:
	types: [opened, synchronize, reopened, labeled]
	branches: [main]

	jobs:
	benchmark-pr:
	name: Benchmark PR
	if: contains(github.event.pull_request.labels.*.name, 'benchmark PR')
	runs-on: ubuntu-latest
	steps:
	- name: Checkout
	uses: actions/checkout@v3
	- name: Setup MPI
	uses: mpi4py/setup-mpi@v1
	- name: Use Python 3.10
	uses: actions/setup-python@v4
	with:
	python-version: 3.10.11 # Perun only supports 3.8 and ahead
	architecture: x64
	- name: Test
	run: \|
	pip install torch==1.12.1+cpu torchvision==0.13.1+cpu torchaudio==0.12.1 -f https://download.pytorch.org/whl/torch_stable.html
	pip install .[cb]
	PERUN_RUN_ID=N4 mpirun -n 4 python benchmarks/cb/main.py
	jq -s flatten bench_data/*.json > bench_data/all_benchmarks.json
	- name: Compare benchmark result
	if: ${{github.ref != 'refs/heads/main'}}
	uses: benchmark-action/github-action-benchmark@v1
	with:
	github-token: ${{secrets.GITHUB_TOKEN}}
	# Benchmark action input and output
	tool: 'customSmallerIsBetter'
	output-file-path: bench_data/all_benchmarks.json
	# external-data-json-path: ./cache/benchmark-data.json
	# Alert configuration
	fail-on-alert: true
	comment-on-alert: true
	# Ignore results from non main branches.
	save-data-file: false
	# Pages configuration
	auto-push: false
	gh-pages-branch: gh-pages
	benchmark-data-dir-path: dev/bench

Provide feedback