Revert "Fix JAX install method (#2485)"

README.md

@@ -49,9 +49,8 @@ Intel Neural Compressor supports PyTorch with CPU, GPU and HPU. Please install t
 pip install neural-compressor-pt
 # Framework extension API + TensorFlow dependency
 pip install neural-compressor-tf
-# Framework extension API + JAX dependency, available since v3.8
-# JAX only support build from source installation method before [PyPI support](https://github.com/pypi/support/issues/10012) is available
-INC_JAX_ONLY=1 pip install . 
+# Framework extension API + JAX dependency
+pip install neural-compressor-jax
 ```    
 **Note**: Further installation methods can be found under [Installation Guide](./docs/source/installation_guide.md). check out our [FAQ](./docs/source/faq.md) for more details.
 

docs/source/installation_guide.md

@@ -57,6 +57,10 @@ pip install neural-compressor-pt
 # Framework extension API + TensorFlow dependency
 pip install neural-compressor-tf
 ```
+```Shell
+# Framework extension API + JAX dependency, available since v3.9
+pip install neural-compressor-jax
+```
 
 ### Install from Source
 The latest code on main branch may not be stable. Please switch to the latest release tag for better stability. Feel free to open an [issue](https://github.com/intel/neural-compressor/issues) if you encounter an error.  

examples/jax/keras/gemma/README.md

@@ -9,7 +9,12 @@ This document describes quantization of Keras Gemma models using Neural Compress
 It is worth conducting experiments in a separate environment. For example, you can use the conda environment from [conda-forge](https://github.com/conda-forge/miniforge). The binary for your environment could be found here: [miniforge](https://github.com/conda-forge/miniforge/releases/latest)  
 
 ## 2. Install modules
-Install Neural Compressor from the source code:
+
+Since Neural Compressor v3.9, you can install the required binaries from PyPI:
+```bash
+pip install -r requirements.txt
+```
+Alternatively, you can always install Neural Compressor from the source code:
 ```bash
 pushd ../../../..  # go to the root directory of the Neural Compressor source code
 INC_JAX_ONLY=1 pip install .

examples/jax/keras/gemma/requirements.txt

@@ -0,0 +1 @@
+neural-compressor-jax

examples/jax/keras/vit/README.md

@@ -10,11 +10,11 @@ It is worth conducting experiments in a separate environment. For example, you c
 
 ## 2. Install modules
 
-Install requirements:
+Since Neural Compressor v3.9, you can install the required binaries from PyPI:
 ```bash
 pip install -r requirements.txt
 ```
-Install Neural Compressor from the source code:
+Alternatively, you can always install Neural Compressor from the source code:
 ```bash
 pushd ../../../..  # go to the root directory of the Neural Compressor source code
 INC_JAX_ONLY=1 pip install .

examples/jax/keras/vit/requirements.txt

@@ -1 +1,2 @@
+neural-compressor-jax
 tensorflow