Add JAX example to L0 test (triton-inference-server#4922)

qa/L0_backend_python/examples/test.sh

@@ -37,11 +37,14 @@ SERVER_LOG="./inference_server.log"
 RET=0
 rm -fr *.log python_backend/
 
-# # Skip torch install on Jetson since it is already installed.
+# Install torch
+# Skip torch install on Jetson since it is already installed.
 if [ "$TEST_JETSON" == "0" ]; then
     pip3 uninstall -y torch
     pip3 install torch==1.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
 fi
+# Install JAX
+pip3 install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 git clone https://github.com/triton-inference-server/python_backend -b $PYTHON_BACKEND_REPO_TAG
 cd python_backend
@@ -108,6 +111,39 @@ wait $SERVER_PID
 
 # Example 3
 
+# JAX AddSub
+CLIENT_LOG="./jax_client.log"
+mkdir -p models/jax/1/
+cp examples/jax/model.py models/jax/1/model.py
+cp examples/jax/config.pbtxt models/jax/config.pbtxt
+run_server
+if [ "$SERVER_PID" == "0" ]; then
+    echo -e "\n***\n*** Failed to start $SERVER\n***"
+    cat $SERVER_LOG
+    RET=1
+fi
+
+set +e
+python3 examples/jax/client.py > $CLIENT_LOG
+if [ $? -ne 0 ]; then
+    echo -e "\n***\n*** Failed to verify jax example. \n***"
+    RET=1
+fi
+
+grep "PASS" $CLIENT_LOG
+if [ $? -ne 0 ]; then
+    echo -e "\n***\n*** Failed to verify jax example. \n***"
+    cat $CLIENT_LOG
+    RET=1
+fi
+set -e
+
+kill $SERVER_PID
+wait $SERVER_PID
+
+
+# Example 4
+
 # BLS Sync
 CLIENT_LOG="./sync_client.log"
 mkdir -p models/bls_sync/1
@@ -138,7 +174,7 @@ set -e
 kill $SERVER_PID
 wait $SERVER_PID
 
-# Example 4
+# Example 5
 
 # Decoupled Repeat
 CLIENT_LOG="./repeat_client.log"
@@ -170,7 +206,7 @@ set -e
 kill $SERVER_PID
 wait $SERVER_PID
 
-# Example 5
+# Example 6
 
 # Decoupled Square
 CLIENT_LOG="./square_client.log"