[JAX] Support various implementations of RoPE. #3173

	# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	#
	# See LICENSE for license information.

	# A workflow to trigger lint tests on GitHub
	name: 'Lint'
	on:
	pull_request:
	workflow_dispatch:
	jobs:
	pytorch_cpplint:
	name: 'PyTorch C++'
	runs-on: ubuntu-latest
	steps:
	- name: Checkout
	uses: actions/checkout@v3
	- name: 'Lint'
	run: \|
	export CPP_ONLY=1
	export TE_PATH=.
	bash ./qa/L0_pytorch_lint/test.sh
	pytorch_pylint:
	name: 'PyTorch Python'
	runs-on: ubuntu-latest
	if: false # NGC PyTorch container does not fit on GitHub runner
	container:
	image: nvcr.io/nvidia/pytorch:23.03-py3
	options: --user root
	steps:
	- name: 'Checkout'
	uses: actions/checkout@v3
	- name: 'Lint'
	run: \|
	pip install flash-attn==1.0.2
	export PYTHON_ONLY=1
	export TE_PATH=.
	bash ./qa/L0_pytorch_lint/test.sh
	jax_cpplint:
	name: 'JAX C++'
	runs-on: ubuntu-latest
	steps:
	- name: 'Checkout'
	uses: actions/checkout@v3
	- name: 'Lint'
	run: \|
	export CPP_ONLY=1
	export TE_PATH=.
	bash ./qa/L0_jax_lint/test.sh
	jax_pylint:
	name: 'JAX Python'
	runs-on: ubuntu-latest
	container:
	image: ghcr.io/nvidia/jax:latest
	options: --user root
	steps:
	- name: 'Checkout'
	uses: actions/checkout@v3
	- name: 'Lint'
	run: \|
	export PYTHON_ONLY=1
	export TE_PATH=.
	bash ./qa/L0_jax_lint/test.sh

Provide feedback