Remove TF support (#467)

Signed-off-by: Kirthi Shankar Sivamani <[email protected]>

.github/workflows/build.yml

@@ -47,20 +47,3 @@ jobs:
           NVTE_FRAMEWORK: jax
       - name: 'Sanity check'
         run: python tests/jax/test_sanity_import.py
-  tensorflow:
-    name: 'TensorFlow'
-    runs-on: ubuntu-latest
-    container:
-      image: nvcr.io/nvidia/tensorflow:23.03-tf2-py3
-      options: --user root
-    steps:
-      - name: 'Checkout'
-        uses: actions/checkout@v3
-        with:
-          submodules: recursive
-      - name: 'Build'
-        run: pip install . -v
-        env:
-          NVTE_FRAMEWORK: tensorflow
-      - name: 'Sanity check'
-        run: python tests/tensorflow/test_sanity_import.py

.github/workflows/lint.yml

@@ -63,28 +63,3 @@ jobs:
           export PYTHON_ONLY=1
           export TE_PATH=.
           bash ./qa/L0_jax_lint/test.sh
-  tensorflow_cpplint:
-    name: 'TensorFlow 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_tensorflow_lint/test.sh
-  tensorflow_pylint:
-    name: 'TensorFlow Python'
-    runs-on: ubuntu-latest
-    container:
-      image: nvcr.io/nvidia/tensorflow:23.03-tf2-py3
-      options: --user root
-    steps:
-      - name: 'Checkout'
-        uses: actions/checkout@v3
-      - name: 'Lint'
-        run: |
-          export PYTHON_ONLY=1
-          export TE_PATH=.
-          bash ./qa/L0_tensorflow_lint/test.sh

README.rst

@@ -129,37 +129,6 @@ Flax
         # Update FP8 metas
         other_variables = te.update_fp8_metas(other_grads)
 
-TensorFlow
-^^^^^^^^^^
-
-.. code-block:: python
-
-  import tensorflow as tf
-  import transformer_engine.tensorflow as te
-  from transformer_engine.common import recipe
-  
-  # Set dimensions.
-  in_features = 768
-  out_features = 3072
-  hidden_size = 2048
-  
-  # Initialize model and inputs.
-  model = te.Dense(out_features, use_bias=True)
-  inp = tf.random.normal((hidden_size, in_features))
-  
-  optimizer = tf.keras.optimizers.Adam(0.001)
-  
-  # Create FP8 recipe. Note: All input args are optional.
-  fp8_recipe = recipe.DelayedScaling(margin=0, interval=1, fp8_format=recipe.Format.E4M3)
-  
-  with tf.GradientTape(persistent=True) as tape:
-      # Enables autocasting for the forward pass
-      with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
-          out = model(inp)
-      loss = tf.reduce_sum(out)
-  grads = tape.gradient(loss, model.trainable_variables)
-  optimizer.apply_gradients(zip(grads, model.trainable_variables))
-
 .. overview-end-marker-do-not-remove
 
 Installation

docs/installation.rst

@@ -36,7 +36,6 @@ Additional Prerequisites
 
 1. [For PyTorch support] `PyTorch <https://pytorch.org/>`__ with GPU support.
 2. [For JAX support] `JAX <https://github.com/google/jax/>`__ with GPU support, version >= 0.4.7.
-3. [For TensorFlow support] `TensorFlow <https://www.tensorflow.org/>`__ with GPU support.
 
 Installation (stable release)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -47,7 +46,7 @@ Execute the following command to install the latest stable version of Transforme
 
   pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable
 
-This will automatically detect if any supported deep learning frameworks are installed and build Transformer Engine support for them. To explicitly specify frameworks, set the environment variable `NVTE_FRAMEWORK` to a comma-separated list (e.g. `NVTE_FRAMEWORK=jax,tensorflow`).
+This will automatically detect if any supported deep learning frameworks are installed and build Transformer Engine support for them. To explicitly specify frameworks, set the environment variable `NVTE_FRAMEWORK` to a comma-separated list (e.g. `NVTE_FRAMEWORK=jax,pytorch`).
 
 Installation (development build)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -64,7 +63,7 @@ Execute the following command to install the latest development build of Transfo
 
   pip install git+https://github.com/NVIDIA/TransformerEngine.git@main
 
-This will automatically detect if any supported deep learning frameworks are installed and build Transformer Engine support for them. To explicitly specify frameworks, set the environment variable `NVTE_FRAMEWORK` to a comma-separated list (e.g. `NVTE_FRAMEWORK=jax,tensorflow`).
+This will automatically detect if any supported deep learning frameworks are installed and build Transformer Engine support for them. To explicitly specify frameworks, set the environment variable `NVTE_FRAMEWORK` to a comma-separated list (e.g. `NVTE_FRAMEWORK=jax,pytorch`).
 
 Installation (from source)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^