Repository: allenai/tango (language: python)

Issue:
FLAX Integration
## Motivation

We have had a lot of interest in being able to use TPUs. While PyTorch XLA is an option, the community is more geared towards using JAX/FLAX-based frameworks such as T5X, since the code remains hardware-agnostic. As T5X is a codebase primarily written for the sequential architectures, it can be limiting for researchers. We would like to provide the functionality to experiment with writing custom FLAX-based code within Tango easily.

Note: 

JAX - is a scientific computation library (can be thought of as equivalent to numpy).
FLAX - is a neural network library built using JAX (can be thought of as equivalent to PyTorch).
T5X - is a research framework for sequential models (eg. T5), built using FLAX.

Success criteria:

- [ ] A FLAX training step, along with `Registrable` versions of various functions required (eg. [optimizers]([https://github.com/google/flax/tree/main/flax/optim](https://github.com/google/flax/tree/main/flax/optim)), [lr_schedulers]([https://github.com/google/flax/blob/main/flax/training/lr_schedule.py](https://github.com/google/flax/blob/main/flax/training/lr_schedule.py)), etc.
- [ ] An example Tango config for training one (or more) of the architectures here: [https://github.com/huggingface/transformers/tree/main/examples/flax](https://github.com/huggingface/transformers/tree/main/examples/flax).
- [ ] Benchmarks comparing performances of the same model in PyTorch vs FLAX on different hardware (CPU, GPU, TPU).


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 2815d9117..bf4cbfe10 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -92,7 +92,7 @@ jobs:
               pytest -v --color=yes --doctest-modules tango/integrations/pytorch_lightning tests/integrations/pytorch_lightning
 
           - name: Transformers integration
-            extras: dev,transformers
+            extras: dev,flax,transformers
             requires_torch: true
             run: |
               pytest -v --color=yes --doctest-modules tango/integrations/transformers tests/integrations/transformers
@@ -104,7 +104,7 @@ jobs:
               pytest -v --color=yes --doctest-modules tango/integrations/fairscale tests/integrations/fairscale
 
           - name: W&B integration
-            extras: dev,torch,wandb
+            extras: dev,torch,flax,wandb
             requires_torch: true
             run: |
               pytest -v --color=yes --doctest-modules tango/integrations/wandb tests/integrations/wandb
@@ -115,6 +115,12 @@ jobs:
             run: |
               pytest -v --color=yes --doctest-modules tango/integrations/beaker tests/integrations/beaker
 
+          - name: Flax integration
+            extras: dev,flax,transformers
+            requires_torch: false
+            run: |
+              pytest -v --color=yes --doctest-modules tango/integrations/flax tests/integrations/flax
+
           - name: Example - train_lm
             extras: dev,all
             requires_torch: true
@@ -188,6 +194,12 @@ jobs:
           . .venv/bin/activate
           pip install torch==1.12.0 --extra-index-url https://download.pytorch.org/whl/cpu
 
+      - name: Pre-install flax
+        if: steps.virtualenv-cache.outputs.cache-hit != 'true' && contains(matrix.task.extras, 'flax')
+        run: |
+          . .venv/bin/activate
+          pip install flax==0.5.0 jax==0.3.13 jaxlib==0.3.10 tensorflow-cpu==2.9.1 optax==0.1.3
+
       - name: Install editable (no cache hit)
         if: steps.virtualenv-cache.outputs.cache-hit != 'true'
         run: |
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 29f4e5347..fb9d02a34 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Added Flax Integration along with an example config.
 - Renamed `workspace` parameter of `BeakerWorkspace` class to `beaker_workspace`.
 - `Executor` class is now a `Registrable` base class. `MulticoreExecutor` is registered as "multicore".
 
diff --git a/docs/source/api/integrations/flax.rst b/docs/source/api/integrations/flax.rst
new file mode 100644
index 000000000..4ba1f2092
--- /dev/null
+++ b/docs/source/api/integrations/flax.rst
@@ -0,0 +1,62 @@
+Flax
+=======
+
+.. automodule:: tango.integrations.flax
+
+Reference
+---------
+
+Train step
+~~~~~~~~~~
+
+.. autoclass:: tango.integrations.flax.FlaxTrainStep
+   :members:
+
+.. autoclass:: tango.integrations.flax.TrainConfig
+   :members:
+
+Eval step
+~~~~~~~~~
+
+.. autoclass:: tango.integrations.flax.FlaxEvalStep
+   :members:
+
+Flax format
+~~~~~~~~~~~~
+
+.. autoclass:: tango.integrations.flax.FlaxFormat
+
+Model
+~~~~~
+
+.. autoclass:: tango.integrations.flax.Model
+   :members:
+
+Optim
+~~~~~
+
+.. autoclass:: tango.integrations.flax.Optimizer
+   :members:
+
+.. autoclass:: tango.integrations.flax.LRScheduler
+   :members:
+
+Data
+~~~~
+
+.. autoclass:: tango.integrations.flax.DataLoader
+   :members:
+
+.. autoclass:: tango.integrations.flax.FlaxDataLoader
+   :members:
+
+Callbacks
+~~~~~~~~~
+
+.. autoclass:: tango.integrations.flax.TrainCallback
+   :members:
+   :member-order: bysource
+
+.. autoclass:: tango.integrations.flax.EvalCallback
+   :members:
+   :member-order: bysource
diff --git a/docs/source/api/integrations/index.rst b/docs/source/api/integrations/index.rst
index f44e52159..b0e2952bd 100644
--- a/docs/source/api/integrations/index.rst
+++ b/docs/source/api/integrations/index.rst
@@ -14,3 +14,4 @@ Integrations
    wandb
    pytorch_lightning
    beaker
+   flax
diff --git a/docs/source/api/integrations/wandb.rst b/docs/source/api/integrations/wandb.rst
index 680fe3762..c894e2bb6 100644
--- a/docs/source/api/integrations/wandb.rst
+++ b/docs/source/api/integrations/wandb.rst
@@ -11,3 +11,5 @@ Reference
 .. autoclass:: tango.integrations.wandb.WandbStepCache
 
 .. autoclass:: tango.integrations.wandb.WandbTrainCallback
+
+.. autoclass:: tango.integrations.wandb.WandbFlaxTrainCallback
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 3c7b543f9..c9526bb8c 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -63,6 +63,7 @@
     "pytorch_lightning": ("https://pytorch-lightning.readthedocs.io/en/latest", None),
     # Just leaving this here for when we do finally add a deepspeed integration:
     #  "deepspeed": ("https://deepspeed.readthedocs.io/en/latest/", None),
+    "flax": ("https://flax.readthedocs.io/en/latest", None),
     "fairscale": ("https://fairscale.readthedocs.io/en/latest/", None),
     "datasets": ("https://huggingface.co/docs/datasets/master/en", None),
     "transformers": ("https://huggingface.co/docs/transformers/master/en", None),
diff --git a/examples/flax/config.jsonnet b/examples/flax/config.jsonnet
new file mode 100644
index 000000000..f9229d797
--- /dev/null
+++ b/examples/flax/config.jsonnet
@@ -0,0 +1,72 @@
+{
+    "steps": {
+        "data": {
+            "type": "datasets::load",
+            "path": "xsum",
+        },
+        "tokenize": {
+            "type": "tokenize_data",
+            "dataset": {
+                "type": "ref",
+                "ref": "data"
+            }
+        },
+        "train": {
+            "type": "flax::train",
+            "model": {
+                "type" : "transformers::FlaxAutoModelForSeq2SeqLM::from_pretrained",
+                "pretrained_model_name_or_path" : "facebook/bart-base"
+            },
+            "dataset": {
+                "type": "ref",
+                "ref": "tokenize"
+            },
+            "optimizer": {
+                "type" : "optax::adamw",
+                "learning_rate" : 2e-5
+            },
+            "train_dataloader": {
+                "batch_size": 16,
+                "drop_last": true
+            },
+            "wrapper": {
+                "type": "xsum_wrapper"
+            },
+            "train_split": "train",
+            "validation_split" : "validation",
+            "validate_every" : 1000,
+            "validation_dataloader": {
+                "batch_size": 16,
+                "drop_last": true
+            },
+            "train_epoch": 5,
+            "checkpoint_every": 1000,
+            "log_every": 1000,
+
+            "callbacks" : [{
+                "type" : "wandb::log_flax"
+            },
+            {
+                "type": "flax::generate_step"
+            }]
+        },
+        "eval": {
+            "type": "flax::eval",
+            "state": {
+                "type": "ref",
+                "ref": "train"
+            },
+            "dataset": {
+                "type": "ref",
+                "ref" : "tokenize"
+            },
+            "dataloader": {
+                "batch_size": 16,
+                "drop_last": true
+            },
+            "wrapper": {
+                "type" : "xsum_wrapper"
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/examples/flax/run.sh b/examples/flax/run.sh
new file mode 100755
index 000000000..22183d611
--- /dev/null
+++ b/examples/flax/run.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+tango run config.jsonnet -d workspace --include-package xsum
diff --git a/examples/flax/xsum.py b/examples/flax/xsum.py
new file mode 100644
index 000000000..1c93dd522
--- /dev/null
+++ b/examples/flax/xsum.py
@@ -0,0 +1,189 @@
+import logging
+from typing import List, Optional
+
+import jax
+import jax.numpy as jnp
+import nltk
+import numpy as np
+import optax
+from datasets import load_metric
+from flax.training.common_utils import onehot
+from transformers import AutoConfig, AutoTokenizer, FlaxAutoModelForSeq2SeqLM
+
+from tango.integrations.flax import FlaxWrapper
+from tango.integrations.flax.train_callback import TrainCallback
+from tango.step import Step
+
+"""
+XSum Summarization with facebook/bart-base
+"""
+
+
+@Step.register("tokenize_data")
+class PreProcessing(Step):
+    DETERMINISTIC = False
+
+    def run(self, dataset):
+        tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
+        model = FlaxAutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
+        model_module = __import__(model.__module__, fromlist=["shift_tokens_tight"])
+        shift_tokens_right_fn = getattr(model_module, "shift_tokens_right")
+        config = AutoConfig.from_pretrained("facebook/bart-base")
+
+        MAX_SOURCE_LENGTH = 512
+        MAX_TGT_LENGTH = 64
+
+        def preprocess_function(examples):
+            inputs = examples["document"]
+            targets = examples["summary"]
+            inputs = [inp for inp in inputs]
+            model_inputs = tokenizer(
+                inputs,
+                max_length=MAX_SOURCE_LENGTH,
+                padding="max_length",
+                truncation=True,
+                return_tensors="np",
+            )
+
+            # Setup the tokenizer for targets
+            with tokenizer.as_target_tokenizer():
+                labels = tokenizer(
+                    targets,
+                    max_length=MAX_TGT_LENGTH,
+                    padding="max_length",
+                    truncation=True,
+                    return_tensors="np",
+                )
+
+            model_inputs["labels"] = labels["input_ids"]
+            decoder_input_ids = shift_tokens_right_fn(
+                labels["input_ids"], config.pad_token_id, config.decoder_start_token_id
+            )
+            model_inputs["decoder_input_ids"] = np.asarray(decoder_input_ids)
+
+            # We need decoder_attention_mask so we can ignore pad tokens from loss
+            model_inputs["decoder_attention_mask"] = labels["attention_mask"]
+
+            return model_inputs
+
+        column_names = dataset["train"].column_names
+
+        dataset = dataset.map(
+            preprocess_function,
+            batched=True,
+            remove_columns=column_names,
+            desc="Running tokenizer on dataset",
+        )
+
+        return dataset
+
+
+@FlaxWrapper.register("xsum_wrapper")  # type: ignore
+class TransformerWrapper(FlaxWrapper):
+    def compute_metrics(self, state, batch, labels):
+        # return empty dict if no other metrics to compute
+        return {}
+
+    def loss_helper(self, logits, labels, batch):
+        label_smoothing_factor = 0
+        padding_mask = batch["decoder_attention_mask"]
+        vocab_size = logits.shape[-1]
+        confidence = 1.0 - label_smoothing_factor
+        low_confidence = (1.0 - confidence) / (vocab_size - 1)
+        normalizing_constant = -(
+            confidence * jnp.log(confidence)
+            + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)
+        )
+        soft_labels = onehot(labels, vocab_size, on_value=confidence, off_value=low_confidence)
+
+        loss = optax.softmax_cross_entropy(logits, soft_labels)
+        loss = loss - normalizing_constant
+
+        # ignore padded tokens from loss
+        loss = loss * padding_mask
+        loss = loss.sum() / padding_mask.sum()
+        return loss
+
+    def train_loss(self, params, state, batch, dropout_rng, labels):
+        logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
+        loss = self.loss_helper(logits, labels, batch)
+        return loss
+
+    def val_metrics(self, batch, logits, labels):
+        loss = self.loss_helper(logits, labels, batch)
+        metrics = {"loss": loss}
+        return metrics
+
+    def eval_metrics(self, batch, logits, labels):
+        loss = self.loss_helper(logits, labels, batch)
+        metrics = {"loss": loss}
+        return metrics
+
+
+@TrainCallback.register("flax::generate_step")
+class GenerateCallback(TrainCallback):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.logger = logging.getLogger(GenerateCallback.__name__)
+
+    def generate_step(self, params, batch):
+        self.model.params = params
+        gen_kwargs = {"max_length": 64, "num_beams": self.model.config.num_beams}
+        output_ids = self.model.generate(
+            batch["input_ids"], attention_mask=batch["attention_mask"], **gen_kwargs
+        )
+        return output_ids.sequences
+
+    def pre_train_loop(self) -> None:
+        if len(jax.devices()) > 1:
+            self.p_generate_step = jax.pmap(self.generate_step, axis_name="batch")
+
+    def pre_val_loop(self, step: int, val_step: int, state) -> None:
+        self.state = state
+        self.eval_preds: List = []
+        self.eval_labels: List = []
+
+    def pre_val_batch(self, step: int, val_step: int, epoch: int, val_batch) -> None:
+        labels = val_batch["labels"]
+        if len(jax.devices()) > 1:
+            generated_ids = self.p_generate_step(self.state.params, val_batch)
+        else:
+            generated_ids = self.generate_step(self.state.params, val_batch)
+        self.eval_preds.extend(jax.device_get(generated_ids.reshape(-1, 64)))
+        self.eval_labels.extend(jax.device_get(labels.reshape(-1, labels.shape[-1])))
+
+    def postprocess_text(self, preds, labels):
+        preds = [pred.strip() for pred in preds]
+        labels = [label.strip() for label in labels]
+
+        # rougeLSum expects newline after each sentence
+        preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds]
+        labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels]
+
+        return preds, labels
+
+    def compute_metrics(self, preds, labels):
+        tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
+        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
+        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+        # Some simple post-processing
+        decoded_preds, decoded_labels = self.postprocess_text(decoded_preds, decoded_labels)
+        metric = load_metric("rouge")
+        result = metric.compute(
+            predictions=decoded_preds, references=decoded_labels, use_stemmer=True
+        )
+        # Extract a few results from ROUGE
+        result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
+
+        prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
+        result["gen_len"] = np.mean(prediction_lens)
+        result = {k: round(v, 4) for k, v in result.items()}
+        return result
+
+    def post_val_loop(
+        self, step: int, epoch: int, val_metric: Optional[float], best_val_metric: Optional[float]
+    ) -> None:
+        rouge_metrics = self.compute_metrics(self.eval_preds, self.eval_labels)
+        rouge_desc = " ".join([f"Eval {key}: {value} |" for key, value in rouge_metrics.items()])
+        self.logger.info(rouge_desc)
diff --git a/requirements.txt b/requirements.txt
index 949112ab1..1059fa27b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -39,4 +39,7 @@ beaker-py>=1.6.2,<2.0.0      # needed by: beaker
 
 # sacremoses should be a dependency of transformers, but it is missing, so we add it manually.
 sacremoses                   # needed by: transformers
-
+jax>=0.3.13                  # needed by: flax
+flax>=0.5.0                  # needed by: flax
+optax>=0.1.2                 # needed by: flax
+tensorflow-cpu>=2.9.1        # needed by: flax
\ No newline at end of file
diff --git a/tango/integrations/flax/__init__.py b/tango/integrations/flax/__init__.py
new file mode 100644
index 000000000..d10c5478f
--- /dev/null
+++ b/tango/integrations/flax/__init__.py
@@ -0,0 +1,32 @@
+from tango.common.exceptions import IntegrationMissingError
+
+try:
+    import flax
+except ModuleNotFoundError:
+    raise IntegrationMissingError("flax")
+
+__all__ = [
+    "DataLoader",
+    "FlaxDataLoader",
+    "LRScheduler",
+    "Model",
+    "Optimizer",
+    "FlaxTrainStep",
+    "FlaxFormat",
+    "TrainCallback",
+    "EvalCallback",
+    "FlaxWrapper",
+    "TrainConfig",
+    "FlaxEvalStep",
+]
+
+from .data import DataLoader, FlaxDataLoader
+from .eval import FlaxEvalStep
+from .eval_callback import EvalCallback
+from .format import FlaxFormat
+from .model import Model
+from .optim import LRScheduler, Optimizer
+from .train import FlaxTrainStep
+from .train_callback import TrainCallback
+from .train_config import TrainConfig
+from .wrapper import FlaxWrapper
diff --git a/tango/integrations/flax/data.py b/tango/integrations/flax/data.py
new file mode 100644
index 000000000..be4d8a04d
--- /dev/null
+++ b/tango/integrations/flax/data.py
@@ -0,0 +1,59 @@
+import logging
+from typing import Generic, TypeVar
+
+import jax.random
+import numpy as np
+from flax.training.common_utils import shard
+
+from tango.common.dataset_dict import DatasetDictBase
+from tango.common.registrable import Registrable
+
+T = TypeVar("T")
+
+
+class DataLoader(Generic[T], Registrable):
+    """
+    A :class:`~tango.common.Registrable` version of a ``Flax DataLoader``.
+    ``Flax DataLoader`` accepts Dataset object. The class yields a numpy batch.
+    """
+
+
+@DataLoader.register("flax::dataloader")
+class FlaxDataLoader(DataLoader):
+    def __init__(
+        self,
+        dataset: DatasetDictBase,
+        batch_size: int = 8,
+        drop_last: bool = True,
+        shuffle: bool = True,
+    ):
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.drop_last = drop_last
+        self.shuffle = shuffle
+        self.dataset_size = self._get_size()
+
+        self.logger = logging.getLogger(FlaxDataLoader.__name__)
+
+    def _get_size(self):
+        size = self.dataset["num_rows"] if type(self.dataset) is dict else self.dataset.num_rows
+        return size
+
+    def __call__(self, rng: jax.random.PRNGKeyArray, do_distributed: bool):
+        steps_per_epoch = self.dataset_size // self.batch_size
+
+        if self.shuffle:
+            perms = jax.random.permutation(rng, self.dataset_size)
+            perms = np.asarray(perms)  # using jax arrays for indexing is a bottleneck on TPUs.
+        else:
+            perms = np.arange(self.dataset_size)
+
+        self.logger.info("Skipping last incomplete batch")
+        perms = perms[: steps_per_epoch * self.batch_size]  # Skip incomplete batch.
+        perms = perms.reshape((steps_per_epoch, self.batch_size))
+
+        for perm in perms:
+            batch = self.dataset[perm]
+            if do_distributed:
+                batch = shard(batch)
+            yield batch
diff --git a/tango/integrations/flax/eval.py b/tango/integrations/flax/eval.py
new file mode 100644
index 000000000..5944211a3
--- /dev/null
+++ b/tango/integrations/flax/eval.py
@@ -0,0 +1,205 @@
+import logging
+from collections import defaultdict
+from itertools import islice
+from typing import Dict, List, Optional, Sequence
+
+import jax
+from flax import jax_utils
+from flax.training.train_state import TrainState
+
+from tango.common.dataset_dict import DatasetDictBase
+from tango.common.exceptions import ConfigurationError
+from tango.common.lazy import Lazy
+from tango.common.tqdm import Tqdm
+from tango.format import Format, JsonFormat
+from tango.step import Step
+
+from .data import FlaxDataLoader
+from .eval_callback import EvalCallback
+from .util import get_PRNGkey
+from .wrapper import FlaxWrapper
+
+
+@Step.register("flax::eval")
+class FlaxEvalStep(Step):
+    """
+    A Flax evaluation loop that pairs well with :class:`FlaxTrainStep`.
+
+    .. tip::
+
+        Registered as a :class:`~tango.step.Step` under the name "flax::eval".
+
+    .. important::
+
+        The evaluation loop will use a GPU/TPU automatically if one is available.
+        You can control which GPU it uses with the environment variable ``CUDA_VISIBLE_DEVICES``.
+        For example, set ``CUDA_VISIBLE_DEVICES=1`` to force ``FlaxEvalStep`` to only use
+        the GPU with ID 1.
+
+    .. warning::
+
+        By default the metrics specified by the ``metric_names`` parameter
+        are aggregated by simply averaging across batches.
+        This behavior is usually correct for metrics like "loss" or "accuracy",
+        for example, but may not be correct for other metrics like "F1".
+
+        If this is not correct for your metric you will need to handle the aggregation
+        internally in your model or with an :class:`EvalCallback`
+        using the :meth:`EvalCallback.post_batch()` method.
+        Then set the parameter ``auto_aggregate_metrics`` to ``False``.
+
+    """
+
+    DETERMINISTIC = True
+    CACHEABLE = True
+    FORMAT: Format = JsonFormat()
+    SKIP_ID_ARGUMENTS = {"log_every"}
+
+    def run(  # type: ignore[override]
+        self,
+        state: TrainState,
+        dataset: DatasetDictBase,
+        dataloader: Lazy[FlaxDataLoader],
+        wrapper: FlaxWrapper,
+        test_split: str = "test",
+        seed: int = 42,
+        log_every: int = 1,
+        do_distributed: bool = False,
+        eval_steps: Optional[int] = None,
+        metric_names: Sequence[str] = ("loss",),
+        auto_aggregate_metrics: bool = True,
+        callbacks: Optional[List[Lazy[EvalCallback]]] = None,
+    ) -> Dict[str, float]:
+        """
+        Evaluate the ``model``.
+
+        :param state:
+            The state of the model to evaluate. This contains the parameters.
+        :param dataset:
+            Should contain the test data.
+        :param dataloader:
+            The data loader that generates test batches. The batches should be :class:`dict`
+            objects.
+        :param wrapper:
+            The wrapper should define :meth:`eval_metrics`.
+        :param test_split:
+            The name of the data split used for evaluation in the ``dataset_dict``.
+            Default is "test".
+        :param seed:
+             Used to set the PRNG states at the beginning of the evaluation loop.
+        :param log_every:
+            Log every this many steps. Default is ``1``.
+        :param do_distributed:
+            Whether to do distributed training or not. Set as 0 or 1.
+        :param eval_steps:
+            The number of steps to evaluate for. If not specified evaluation will
+            stop after a complete iteration through the ``dataloader``.
+        :param metric_names:
+            The names of the metrics to track and aggregate. Default is ``("loss",)``.
+        :param auto_aggregate_metrics:
+            If ``True`` (the default), the metrics will be averaged across batches.
+            This may not be the correct behavior for some metrics (such as F1),
+            in which you should set this to ``False`` and handle the aggregation
+            internally in your model or with an :class:`EvalCallback`
+            (using :meth:`EvalCallback.post_batch()`).
+        :param callbacks:
+            A list of :class:`EvalCallback`.
+
+        """
+
+        logger = logging.getLogger(FlaxEvalStep.__name__)
+        # construct dataloader
+        dataloader: FlaxDataLoader = dataloader.construct(
+            dataset=dataset[test_split].set_format("numpy")
+        )
+
+        steps: int
+        try:
+            dataloader_len = dataloader.dataset_size
+            steps = dataloader_len if eval_steps is None else min(dataloader_len, eval_steps)
+        except TypeError:
+            if eval_steps is None:
+                raise ConfigurationError(
+                    "You must set 'eval_steps' for streaming/iterable datasets"
+                )
+            else:
+                steps = eval_steps
+
+        if do_distributed:
+            devices = jax.devices()
+            if len(devices) <= 1:
+                raise ConfigurationError(
+                    "YOu have set distributed training=True but there is only one device."
+                )
+
+        # Initialize callbacks
+        callbacks: List[EvalCallback] = [
+            callback.construct(
+                step_id=self.unique_id,
+                work_dir=self.work_dir,
+                dataset_dict=dataset,
+                dataloader=dataloader,
+            )
+            for callback in (callbacks or [])
+        ]
+
+        for callback in callbacks:
+            callback.pre_eval_loop()
+
+        rng = get_PRNGkey(seed)
+        devices = jax.devices()
+        if len(devices) > 1:
+            do_distributed = True
+
+        def eval_step(state, batch):
+            labels = batch.pop("labels")
+            logits = state.apply_fn(**batch, params=state.params, train=False)[0]
+            metrics = wrapper.eval_metrics(batch, logits, labels)
+            if do_distributed:
+                metrics = jax.lax.pmean(metrics, axis_name="batch")
+            return logits, metrics
+
+        if do_distributed:
+            state = jax_utils.replicate(state)
+            parallel_eval_step = jax.pmap(eval_step, axis_name="batch")
+
+        eval_batches = enumerate(islice(dataloader(rng, do_distributed), steps))
+
+        running_metrics: Dict[str, float] = defaultdict(float)
+        aggregated_metrics: Dict[str, float] = defaultdict(float)
+
+        with Tqdm.tqdm(eval_batches, desc="Evaluating", total=steps) as batch_iter:
+            for step, batch in batch_iter:
+                should_log_this_step = step % log_every == 0 or step == steps - 1
+                for callback in callbacks:
+                    callback.pre_batch(step, batch)
+
+                if do_distributed:
+                    logits, metrics = parallel_eval_step(state, batch)
+                    metrics = jax_utils.unreplicate(metrics)
+                else:
+                    logits, metrics = eval_step(state, batch)
+
+                for callback in callbacks:
+                    callback.post_batch(step, logits)
+
+                if auto_aggregate_metrics:
+                    for key, val in metrics.items():
+                        if key in metric_names:
+                            running_metrics[key] += metrics[key].item()
+                            aggregated_metrics[key] = running_metrics[key] / (step + 1)
+                else:
+                    aggregated_metrics.update(metrics)
+
+                if should_log_this_step:
+                    batch_iter.set_postfix(**aggregated_metrics)
+                del batch
+
+        logger.info("Evaluation Metrics:")
+        for key, val in aggregated_metrics.items():
+            logger.info(key, ":", val)
+
+        for callback in callbacks:
+            callback.post_eval_loop(aggregated_metrics)
+
+        return aggregated_metrics
diff --git a/tango/integrations/flax/eval_callback.py b/tango/integrations/flax/eval_callback.py
new file mode 100644
index 000000000..e636127ab
--- /dev/null
+++ b/tango/integrations/flax/eval_callback.py
@@ -0,0 +1,74 @@
+from pathlib import Path
+from typing import Any, Dict
+
+from tango.common.dataset_dict import DatasetDictBase
+from tango.common.registrable import Registrable
+from tango.workspace import Workspace
+
+from .data import FlaxDataLoader
+
+
+class EvalCallback(Registrable):
+    """
+    An ``EvalCallback`` is a :class:`~tango.common.Registrable` class that can be used
+    within :class:`FlaxEvalStep` to customize the behavior of the evaluation loop,
+    similar to how :class:`TrainCallback` is used to customize the behavior of the training
+    loop.
+
+    .. tip::
+        All of the parameters to this base class will be automatically set within
+        the training loop, so you shouldn't include them in your config for your callbacks.
+
+    :ivar Workspace workspace: The tango workspace being used.
+    :ivar str step_id: The unique ID of the step.
+    :ivar pathlib.Path work_dir: The working directory of the step
+    :ivar DatasetDictBase dataset_dict: The dataset dict containing the evaluation split.
+    :ivar DataLoader dataloader: The data loader used to load the evaluation split data.
+    """
+
+    def __init__(
+        self,
+        workspace: Workspace,
+        step_id: str,
+        work_dir: Path,
+        dataset_dict: DatasetDictBase,
+        dataloader: FlaxDataLoader,
+    ) -> None:
+        self.workspace = workspace
+        self.step_id = step_id
+        self.work_dir = work_dir
+        self.dataset_dict = dataset_dict
+        self.dataloader = dataloader
+
+    def pre_eval_loop(self) -> None:
+        """
+        Called right before the first batch is processed
+        """
+        pass
+
+    def post_eval_loop(self, aggregated_metrics: Dict[str, float]) -> None:
+        """
+        Called after the evaluation loop completes with the final aggregated metrics.
+
+        This is the last method that is called, so any cleanup can be done in this method.
+        """
+        pass
+
+    def pre_batch(self, step: int, batch: Dict[str, Any]) -> None:
+        """
+        Called directly before processing a batch.
+        """
+        pass
+
+    def post_batch(self, step: int, batch_outputs: Dict[str, Any]) -> None:
+        """
+        Called directly after processing a batch with the outputs of the batch.
+
+        .. tip::
+            This method can be used to modify ``batch_outputs`` in place, which is useful
+            in scenarios where you might need to aggregate metrics
+            in a special way other than a simple average. If that's the case, make sure
+            to set ``auto_aggregate_metrics`` to ``False`` in :class:`FlaxEvalStep`.
+
+        """
+        pass
diff --git a/tango/integrations/flax/format.py b/tango/integrations/flax/format.py
new file mode 100644
index 000000000..6a5883916
--- /dev/null
+++ b/tango/integrations/flax/format.py
@@ -0,0 +1,29 @@
+from pathlib import Path
+from typing import Generic, TypeVar
+
+from flax.training import checkpoints
+
+from tango.common.aliases import PathOrStr
+from tango.format import Format
+
+T = TypeVar("T")
+
+
+@Format.register("flax")
+class FlaxFormat(Format[T], Generic[T]):
+    """
+    This format writes the artifact.
+
+    .. tip::
+
+        Registered as a :class:`~tango.format.Format` under the name "flax".
+    """
+
+    VERSION = "002"
+
+    def write(self, artifact: T, dir: PathOrStr) -> None:
+        checkpoints.save_checkpoint(Path(dir), artifact, step=0)
+
+    def read(self, dir: PathOrStr) -> T:
+        # will return a dict
+        return checkpoints.restore_checkpoint(dir, target=None)
diff --git a/tango/integrations/flax/model.py b/tango/integrations/flax/model.py
new file mode 100644
index 000000000..af492a9f0
--- /dev/null
+++ b/tango/integrations/flax/model.py
@@ -0,0 +1,14 @@
+from flax import linen as nn
+
+from tango.common.registrable import Registrable
+
+
+class Model(nn.Module, Registrable):
+    """
+    This is a :class:`~tango.common.Registrable` mixin class that inherits from
+    :class:`flax.linen.nn.Module`.
+    Its :meth:`~flax.linen.nn.Module.setup()` can be used to register submodules,
+    variables, parameters you will need in your model.
+    Its :meth:`~flax.linen.nn.Module.__call__()` returns the output of the model
+    for a given input.
+    """
diff --git a/tango/integrations/flax/optim.py b/tango/integrations/flax/optim.py
new file mode 100644
index 000000000..920991560
--- /dev/null
+++ b/tango/integrations/flax/optim.py
@@ -0,0 +1,109 @@
+from inspect import isfunction
+from typing import Callable, Type
+
+import optax
+
+from tango.common.registrable import Registrable
+
+
+class Optimizer(Registrable):
+    """
+    A :class:`~tango.common.Registrable` version of Optax optimizers.
+
+    All `built-in Optax optimizers
+    <https://optax.readthedocs.io/en/latest/api.html#>`_
+    are registered according to their class name (e.g. "optax::adam").
+
+    .. tip::
+
+        You can see a list of all available optimizers by running
+
+        .. testcode::
+
+            from tango.integrations.flax import Optimizer
+            for name in sorted(Optimizer.list_available()):
+                print(name)
+
+        .. testoutput::
+            :options: +ELLIPSIS
+
+            optax::adabelief
+            optax::adafactor
+            optax::adagrad
+            optax::adam
+            ...
+
+    """
+
+    def __init__(self, optimizer: Callable) -> None:
+        self.optimizer = optimizer
+
+    def __call__(self, **kwargs) -> optax.GradientTransformation:
+        return self.optimizer(**kwargs)
+
+
+class LRScheduler(Registrable):
+    """
+    A :class:`~tango.common.Registrable` version of an Optax learning
+    rate scheduler.
+
+    All `built-in Optax learning rate schedulers
+    <https://optax.readthedocs.io/en/latest/api.html#schedules>`_
+    are registered according to their class name (e.g. "optax::linear_schedule").
+
+    .. tip::
+
+        You can see a list of all available schedulers by running
+
+        .. testcode::
+
+            from tango.integrations.flax import LRScheduler
+            for name in sorted(LRScheduler.list_available()):
+                print(name)
+
+        .. testoutput::
+            :options: +ELLIPSIS
+
+            optax::constant_schedule
+            optax::cosine_decay_schedule
+            optax::cosine_onecycle_schedule
+            optax::exponential_decay
+            ...
+
+    """
+
+    def __init__(self, scheduler: Callable) -> None:
+        self.scheduler = scheduler
+
+    def __call__(self, **kwargs):
+        return self.scheduler(**kwargs)
+
+
+def optimizer_factory(optim_method: Callable) -> Type[Callable]:
+    def factory_func():
+        return Optimizer(optim_method)
+
+    return factory_func()
+
+
+def scheduler_factory(scheduler_method: Callable) -> Type[Callable]:
+    def factory_func():
+        return LRScheduler(scheduler_method)
+
+    return factory_func()
+
+
+# Register all optimizers.
+for name, cls in optax._src.alias.__dict__.items():
+    if isfunction(cls) and not name.startswith("_") and cls.__annotations__:
+        factory_func = optimizer_factory(cls)
+        Optimizer.register("optax::" + name)(factory_func)
+
+# Register all learning rate schedulers.
+for name, cls in optax._src.schedule.__dict__.items():
+    if isfunction(cls) and not name.startswith("_") and cls.__annotations__:
+        factory_func = scheduler_factory(cls)
+        LRScheduler.register("optax::" + name)(factory_func)
+
+# TODO: Handle inject_hyperparams.
+# Refer: https://optax.readthedocs.io/en/latest/api.html?highlight=inject%20hyperparam
diff --git a/tango/integrations/flax/train.py b/tango/integrations/flax/train.py
new file mode 100644
index 000000000..32af99dfd
--- /dev/null
+++ b/tango/integrations/flax/train.py
@@ -0,0 +1,607 @@
+import logging
+import time
+from collections import defaultdict
+from pathlib import Path
+from typing import Any, DefaultDict, Dict, List, Optional
+
+import jax
+import jax.numpy as jnp
+from flax import jax_utils
+from flax.training import checkpoints
+from flax.training.train_state import TrainState
+
+from tango.common.dataset_dict import DatasetDictBase
+from tango.common.exceptions import ConfigurationError
+from tango.common.lazy import Lazy
+from tango.common.tqdm import Tqdm
+from tango.format import Format
+from tango.step import Step
+from tango.workspace import Workspace
+
+from .data import FlaxDataLoader
+from .format import FlaxFormat
+from .model import Model
+from .optim import LRScheduler, Optimizer
+from .train_callback import TrainCallback
+from .train_config import TrainConfig
+from .util import get_multiple_keys, get_PRNGkey
+from .wrapper import FlaxWrapper
+
+PyTree = Any
+
+
+@Step.register("flax::train")
+class FlaxTrainStep(Step):
+    """
+    A Flax training step that supports distributed training with configurable dataloaders, callbacks,
+    optimizer.
+
+    .. tip::
+        Registered as a :class:`~tango.step.Step` under the name "flax::train".
+
+    .. important::
+        To train on GPUs and TPUs, installation of jax[cuda] or jax[tpu] is required. Follow the
+        instructions here: https://github.com/google/jax to set up jax for GPUs and TPUs.
+        Note: CUDA and cuDNN installation is required to run jax on NVidia GPUs. It is recommended to
+        install cuDNN in your conda environment using: ``conda install -c anaconda cudnn``.
+
+        Distributed data parallel training is activated when the ``device_count`` is greater than 1.
+        You can control which CUDA devices to use with the environment variable ``CUDA_VISIBLE_DEVICES``.
+        For example, to only use the GPUs with IDs 0 and 1, set ``CUDA_VISIBLE_DEVICES=0,1``
+        (and ``device_count`` to 2).
+
+    .. warning::
+        During validation, the validation metric (specified by the ``val_metric_name`` parameter)
+        is aggregated by simply averaging across validation batches and distributed processes.
+        This behavior is usually correct when your validation metric is "loss" or "accuracy",
+        for example, but may not be correct for other metrics like "F1".
+        If this is not correct for your metric you will need to handle the aggregation
+        internally in your model or with a :class:`TrainCallback`
+        using the :meth:`TrainCallback.post_val_batch()` method.
+        Then set the parameter ``auto_aggregate_val_metric`` to ``False``.
+
+        Jax pre-allocates 90% of GPU memory. If you run into out-of-memory (OOM) issues, please refer
+        to this: https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html.
+
+    """
+
+    DETERMINISTIC = True
+    CACHEABLE = True
+    FORMAT: Format = FlaxFormat()
+    SKIP_ID_ARGUMENTS = {"log_every"}
+
+    def run(  # type: ignore[override]
+        self,
+        model: Model,
+        dataset: DatasetDictBase,
+        optimizer: Lazy[Optimizer],
+        train_dataloader: Lazy[FlaxDataLoader],
+        *,
+        wrapper: FlaxWrapper,
+        seed: int = 42,
+        keep_checkpoints: int = 5,
+        lr_scheduler: Optional[Lazy[LRScheduler]] = None,
+        train_split: Optional[str] = None,
+        validation_dataloader: Optional[Lazy[FlaxDataLoader]] = None,
+        validation_split: Optional[str] = None,
+        train_steps: Optional[int] = None,
+        train_epoch: Optional[int] = None,
+        validation_steps: Optional[int] = None,
+        log_every: int = 10,
+        checkpoint_every: int = 100,
+        validate_every: Optional[int] = None,
+        val_metric_name: str = "loss",
+        minimize_val_metric: bool = True,
+        auto_aggregate_val_metric: bool = True,
+        callbacks: Optional[List[Lazy[TrainCallback]]] = None,
+        remove_stale_checkpoints: bool = True,
+    ) -> PyTree:
+        """
+        Run a basic training loop to train the ``model``.
+
+        :param model:
+            The flax model to train. It should define ``__call__()``. Defining ``setup()`` is Optional.
+        :param dataset:
+            The train and optional validation dataset.
+        :param optimizer:
+            The name of the optax Optimizer to use for training.
+        :param train_dataloader:
+            The dataloader object that generates training batches.
+        :param wrapper:
+            A Wrapper class that defines ``loss_fn()``, ``eval_fn()`` and ``compute_metrics()``
+        :param seed:
+            Used to set the PRNG state. By default, ``seed=42``
+        :param keep_checkpoints:
+            An integer which denotes how many previous checkpoints should be stored while training.
+            By default, ``keep_checkpoints=5``
+        :param lr_scheduler:
+            The name of the learning rate scheduler.
+        :param train_split:
+            The name of the data split used for training in the ``dataset_dict``.
+            Default is "train".
+        :param validation_dataloader:
+            An optional data loader for generating validation batches. The batches should be
+            :class:`dict` objects. If not specified, but ``validation_split`` is given,
+            the validation ``DataLoader`` will be constructed from the same parameters
+            as the train ``DataLoader``.
+        :param validation_split:
+            Optional name of the validation split in the ``dataset_dict``. Default is ``None``,
+            which means no validation.
+        :param train_steps:
+            The number of steps to train for. If not specified training will
+            stop after a complete iteration through the ``train_dataloader``.
+        :param train_epoch:
+            The number of epochs to train for. You cannot specify ``train_steps`` and ``train_epochs``
+            at the same time.
+        :param validation_steps:
+            The number of steps to validate for. If not specified validation
+            will stop after a complete iteration through the ``validation_dataloader``.
+        :param log_every:
+            Log every this many steps.
+        :param checkpoint_every:
+            Save a checkpoint every this many steps.
+        :param validate_every:
+            Run the validation loop every this many steps.
+        :param val_metric_name:
+            The name of the validation metric, i.e. the key of the metric in the dictionary
+            returned by the forward pass of the model. Default is "loss".
+        :param minimize_val_metric:
+            Whether the validation metric is meant to be minimized (such as the loss).
+            Default is ``True``. When using a metric such as accuracy, you should set
+            this to ``False``.
+        :param auto_aggregate_val_metric:
+            If ``True`` (the default), the validation metric will be averaged across
+            validation batches and distributed processes. This may not be the correct
+            behavior for some metrics (such as F1), in which you should set this to
+            ``False`` and handle the aggregation internally in your model
+            or with a :class:`TrainCallback` (using :meth:`TrainCallback.post_val_batch()`).
+        :param callbacks:
+            A list of :class: `TrainCallback`.
+        :param remove_stale_checkpoints:
+            If ``True`` (the default), stale checkpoints will be removed throughout training so that
+            only the latest and best checkpoints are kept.
+
+        :returns:
+            The trained model with the last checkpoint loaded.
+        """
+
+        return self._train(
+            dataset=dataset,
+            model=model,
+            optimizer=optimizer,
+            train_dataloader=train_dataloader,
+            train_wrapper=wrapper,
+            seed=seed,
+            keep_checkpoints=keep_checkpoints,
+            lr_scheduler=lr_scheduler,
+            train_split=train_split,
+            validation_split=validation_split,
+            validation_dataloader=validation_dataloader,
+            train_steps=train_steps,
+            train_epochs=train_epoch,
+            validation_steps=validation_steps,
+            log_every=log_every,
+            checkpoint_every=checkpoint_every,
+            validate_every=validate_every,
+            val_metric_name=val_metric_name,
+            minimize_val_metric=minimize_val_metric,
+            auto_aggregate_val_metric=auto_aggregate_val_metric,
+            callbacks=callbacks,
+            remove_stale_checkpoints=remove_stale_checkpoints,
+        )
+
+    def _train(
+        self,
+        model: Model,
+        optimizer: Lazy[Optimizer],
+        dataset: DatasetDictBase,
+        train_dataloader: Lazy[FlaxDataLoader],
+        *,
+        train_wrapper: FlaxWrapper,
+        seed: int = 42,
+        keep_checkpoints: int = 5,
+        lr_scheduler: Optional[Lazy[LRScheduler]],
+        train_split: Optional[str] = "train",
+        validation_split: Optional[str] = None,
+        validation_dataloader: Optional[Lazy[FlaxDataLoader]] = None,
+        train_steps: Optional[int] = None,
+        train_epochs: Optional[int] = None,
+        validation_steps: Optional[int] = None,
+        log_every: int = 10,
+        checkpoint_every: int = 100,
+        validate_every: Optional[int] = None,
+        val_metric_name: str = "loss",
+        minimize_val_metric: bool = True,
+        auto_aggregate_val_metric: bool = True,
+        callbacks: Optional[List[Lazy[TrainCallback]]] = None,
+        remove_stale_checkpoints: bool = True,
+    ) -> PyTree:
+
+        if validate_every is not None and validation_split is None:
+            raise ConfigurationError(
+                "You have set a validation interval, but no validation split. "
+                "That's probably unintentional."
+            )
+
+        if (train_steps is not None) and (train_epochs is not None):
+            raise ConfigurationError(
+                "One of 'train_steps' or 'train_epochs' needs to be specified, but not both."
+            )
+
+        if isinstance(dataset, DatasetDictBase) and train_split is None:
+            raise ConfigurationError("Specify the train split for Datasets object.")
+
+        config = TrainConfig(
+            self.unique_id,
+            self.work_dir,
+            step_name=self.name,
+            train_split=train_split,
+            validation_split=validation_split,
+            seed=seed,
+            train_steps=train_steps,
+            train_epochs=train_epochs,
+            log_every=log_every,
+            checkpoint_every=checkpoint_every,
+            validate_every=validate_every,
+            validation_steps=validation_steps,
+            val_metric_name=val_metric_name,
+            minimize_val_metric=minimize_val_metric,
+            auto_aggregate_val_metric=auto_aggregate_val_metric,
+            remove_stale_checkpoints=remove_stale_checkpoints,
+        )
+
+        optimizer = self._construct_optimizer(optimizer)
+
+        lr_scheduler_: Optional[LRScheduler] = None
+        if lr_scheduler is not None:
+            lr_scheduler_ = self._construct_lr_scheduler(lr_scheduler)
+        lr_scheduler = lr_scheduler_
+
+        final_model: Model
+
+        final_model = self.train_helper(
+            self.workspace,
+            config,
+            model,
+            optimizer,
+            keep_checkpoints,
+            lr_scheduler,
+            train_wrapper,
+            dataset,
+            train_dataloader,
+            validation_dataloader,
+            callbacks,
+        )
+        assert final_model is not None
+
+        return final_model
+
+    def train_helper(
+        self,
+        workspace: Workspace,
+        config: TrainConfig,
+        model: Model,
+        optimizer: Optimizer,
+        keep_checkpoints: int,
+        lr_scheduler: Optional[LRScheduler],
+        train_wrapper: FlaxWrapper,
+        dataset: DatasetDictBase,
+        train_dataloader: Lazy[FlaxDataLoader],
+        validation_dataloader: Optional[Lazy[FlaxDataLoader]] = None,
+        callbacks: Optional[List[Lazy[TrainCallback]]] = None,
+    ) -> PyTree:
+
+        logger = logging.getLogger(FlaxTrainStep.__name__)
+
+        # construct data loaders
+        validation_dataloader_: Optional[FlaxDataLoader] = None
+        if config.validation_split is not None:
+            validation_dataset = dataset[config.validation_split]
+            validation_dataset.set_format("numpy")
+            if validation_dataloader is not None:
+                validation_dataloader_ = validation_dataloader.construct(dataset=validation_dataset)
+            else:
+                validation_dataloader_ = train_dataloader.construct(dataset=validation_dataset)
+
+        validation_dataloader = validation_dataloader_
+
+        train_dataset = dataset
+        if config.train_split is not None:
+            train_dataset = dataset[config.train_split]
+            train_dataset.set_format("numpy")  # type:ignore
+        train_dataloader: FlaxDataLoader = train_dataloader.construct(dataset=train_dataset)
+
+        devices = self._get_devices()
+        do_distributed: bool = False
+        if len(devices) > 1:
+            do_distributed = True
+
+        if validation_dataloader is not None:
+            validation_dataloader.batch_size *= len(devices)
+        train_dataloader.batch_size *= len(devices)
+
+        rng = get_PRNGkey(config.seed)
+
+        if hasattr(model, "params"):
+            params = model.params
+        else:
+            # TODO: Find better way to init the shape
+            shape = list(train_dataset["x"].shape)
+            shape[0] = 1
+            x = jnp.ones(shape)
+
+            params = model.init(rng, x)["params"]
+
+        state = TrainState.create(apply_fn=model.__call__, params=params, tx=optimizer)
+
+        initial_state: Optional[Dict[str, Any]] = None
+        if config.state_path.exists():
+            logger.info("Recovering from previous run at %s" % config.state_path)
+            state = self.load_checkpoint(config.state_path, state)
+
+        if config.train_epochs is None:
+            assert config.train_steps is not None
+            try:
+                train_epochs = len(train_dataloader.dataset) // train_dataloader.batch_size
+            except TypeError:
+                raise ConfigurationError(
+                    "You must set train_epochs for streaming/iterable datasets"
+                )
+
+            config.train_epochs = train_epochs
+
+        assert config.train_epochs is not None
+
+        if validation_dataloader is not None:
+            if config.validation_steps is None:
+                try:
+                    config.validation_steps = len(validation_dataloader.dataset)
+                except TypeError:
+                    raise ConfigurationError(
+                        "You must set 'validation_steps' for streaming/iterable datasets"
+                    )
+
+        val_metric: Optional[float] = None
+        best_val_metric: Optional[float] = None
+        start_step: int = 0
+
+        if initial_state is not None:
+            val_metric = initial_state[f"val_{config.val_metric_name}"]
+            best_val_metric = initial_state[f"best_{config.val_metric_name}"]
+            start_step = initial_state["training_epochs"]
+
+        # Initialize callbacks
+        callbacks: List[TrainCallback] = [
+            callback.construct(
+                workspace=workspace,
+                train_config=config,
+                dataset=dataset,
+                train_dataloader=train_dataloader,
+                model=model,
+                optimizer=optimizer,
+                validation_dataloader=validation_dataloader,
+            )
+            for callback in (callbacks or [])
+        ]
+
+        if initial_state:
+            for callback, state in zip(callbacks, initial_state["callbacks"]):
+                callback.load_state_dict(state)
+
+        del initial_state
+
+        if start_step > 0:
+            with Tqdm.tqdm(
+                train_dataloader,
+                desc=f"Catching dataloader upto step {start_step}",
+                total=start_step - 1,
+            ) as batch_iter:
+                for step, batch in enumerate(batch_iter):
+                    del batch
+                    if step >= start_step - 1:
+                        break
+
+        def train_step(state, batch, dropout_rng):
+            # if transformer model
+            labels = batch.pop("labels")
+            dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
+            grad_fn = jax.value_and_grad(train_wrapper.train_loss)
+            loss, grad = grad_fn(state.params, state, batch, dropout_rng, labels)
+            if do_distributed:
+                grad = jax.lax.pmean(grad, "batch")
+            new_state = state.apply_gradients(grads=grad)
+            other_metrics = train_wrapper.compute_metrics(state, batch, labels=labels)
+            metrics = {"loss": loss}
+            metrics.update(other_metrics)
+            if do_distributed:
+                metrics = jax.lax.pmean(metrics, axis_name="batch")
+
+            return new_state, metrics, new_dropout_rng
+
+        def val_step(state, batch):
+            labels = batch.pop("labels")
+            logits = state.apply_fn(**batch, params=state.params, train=False)[0]
+            metrics = train_wrapper.val_metrics(batch, logits, labels)
+            if do_distributed:
+                metrics = jax.lax.pmean(metrics, axis_name="batch")
+            return metrics
+
+        if do_distributed:
+            # NOTE: The trainer currently handles only data parallelism.
+            state = jax_utils.replicate(state)
+            dropout_rngs = get_multiple_keys(rng, jax.local_device_count())
+            parallel_train_step = jax.pmap(train_step, axis_name="batch")
+            parallel_val_step = jax.pmap(val_step, axis_name="batch")
+
+        step_per_epoch = train_dataloader.dataset_size // train_dataloader.batch_size
+        config.train_steps = step_per_epoch * config.train_epochs
+
+        assert config.train_steps is not None  # for mypy
+
+        for callback in callbacks:
+            callback.pre_train_loop()
+
+        logger.info("***** Running training *****")
+        logger.info(f"  Num examples = {train_dataloader.dataset_size}")
+        logger.info(f"  Num Epochs = {config.train_epochs}")
+        logger.info(
+            f"  Total train batch size (w. parallel & distributed) = {train_dataloader.batch_size}"
+        )
+        logger.info(f"  Total optimization steps = {config.train_steps}")
+
+        step = start_step
+
+        epochs = Tqdm.tqdm(
+            range(config.train_epochs), desc=f"Epoch (1/{config.train_epochs})", position=0
+        )
+
+        for epoch in epochs:
+            start = time.time()
+            train_metrics = []
+
+            for callback in callbacks:
+                callback.pre_epoch(step, epoch)
+
+            train_loader = train_dataloader(rng, do_distributed)
+            for _ in Tqdm.tqdm(range(step_per_epoch), desc="Training", position=1):
+                batch = next(train_loader)
+                for callback in callbacks:
+                    callback.pre_batch(step, epoch, batch)
+
+                if do_distributed:
+                    state, train_metric, dropout_rngs = parallel_train_step(
+                        state, batch, dropout_rngs
+                    )
+                else:
+                    state, train_metric, rng = train_step(state, batch, rng)
+
+                train_metrics.append(train_metric)
+
+                for callback in callbacks:
+                    callback.post_batch(step, epoch, train_metric)
+
+                if config.should_log_this_step(step):
+                    for callback in callbacks:
+                        callback.log_batch(step, epoch, train_metric)
+
+                if config.should_checkpoint_this_step(step):
+                    self.save_checkpoint(config.state_path, state, step, keep_checkpoints)
+                step += 1
+
+                # check if we need to do validation
+                if config.validation_split is None:
+                    # If we can't validate, we don't.
+                    should_validate = False
+                elif step == config.train_steps - 1:
+                    # If we're at the end of the training run, we always validate.
+                    should_validate = True
+                elif config.validate_every is not None and step % config.validate_every == 0:
+                    # If validate_every is given, we use that to decide.
+                    should_validate = True
+                else:
+                    # Otherwise, we don't validate.
+                    should_validate = False
+
+                if should_validate:
+                    assert validation_dataloader is not None
+                    assert config.validation_steps is not None
+
+                    val_metrics: DefaultDict = defaultdict(list)
+                    epoch_eval_metrics: DefaultDict = defaultdict(float)
+
+                    val_dataloader = validation_dataloader(rng, do_distributed)
+
+                    valid_step = 0
+                    total_val_steps = len(validation_dataset) // validation_dataloader.batch_size
+                    for callback in callbacks:
+                        callback.pre_val_loop(step, valid_step, state)
+
+                    for _ in Tqdm.tqdm(range(total_val_steps), desc="Evaluating", position=2):
+                        batch = next(val_dataloader)
+                        for callback in callbacks:
+                            callback.pre_val_batch(step, valid_step, epoch, batch)
+
+                        if do_distributed:
+                            metrics = parallel_val_step(state, batch)
+                            metrics = jax_utils.unreplicate(metrics)
+                        else:
+                            metrics = val_step(state, batch)
+
+                        for key, value in metrics.items():
+                            val_metrics[key].append(value.item())
+
+                        for callback in callbacks:
+                            callback.post_val_batch(step, valid_step, epoch, val_metrics)
+
+                        valid_step += 1
+
+                    for key, value in val_metrics.items():
+                        if config.auto_aggregate_val_metric:
+                            epoch_eval_metrics[key] = jax.tree_map(
+                                jnp.mean, jnp.array(value)
+                            ).item()
+                        else:
+                            epoch_eval_metrics[key] = metrics[key].item()
+
+                    for key, value in epoch_eval_metrics.items():
+                        print("Validation %s : %.5f" % (key, value))
+
+                    val_metric = epoch_eval_metrics[config.val_metric_name]
+
+                    assert val_metric is not None
+
+                    if best_val_metric is None:
+                        best_val_metric = val_metric
+                    elif config.minimize_val_metric and val_metric <= best_val_metric:
+                        best_val_metric = val_metric
+                    elif not config.minimize_val_metric and val_metric >= best_val_metric:
+                        best_val_metric = val_metric
+
+                    for callback in callbacks:
+                        callback.post_val_loop(step, epoch, val_metric, best_val_metric)
+
+            if do_distributed:
+                train_metric = jax_utils.unreplicate(train_metric)
+
+            for key, value in train_metric.items():
+                print("Train %s : %.2f" % (key, value))
+
+            for callback in callbacks:
+                callback.post_epoch(step, epoch)
+
+            end = time.time()
+            train_time = (end - start) / 60
+
+            desc = f"Epoch... ({epoch + 1}/{config.train_epochs} | Time taken (mins): {train_time})"
+            epochs.write(desc)
+            epochs.desc = desc
+
+        for callback in callbacks:
+            callback.post_train_loop(step, epoch)
+
+        if do_distributed:
+            state = jax_utils.unreplicate(state)
+        return state
+
+    def save_checkpoint(self, dir: Path, target: PyTree, step: int, keep_checkpoints: int):
+        return checkpoints.save_checkpoint(
+            dir, target, step, prefix="checkpoint_", keep=keep_checkpoints, overwrite=True
+        )
+
+    def load_checkpoint(self, dir: Path, target: PyTree):
+        return checkpoints.restore_checkpoint(dir, target, prefix="checkpoint_")
+
+    def _construct_optimizer(self, optimizer):
+        self.optimizer = optimizer.construct()
+        return self.optimizer
+
+    def _construct_lr_scheduler(self, scheduler):
+        self.lr_scheduler = scheduler.construct()
+        return self.lr_scheduler
+
+    def _get_devices(self) -> List[Any]:
+        device_type = jax.default_backend()
+        self.devices = jax.devices()
+        device_count = len(self.devices)
+        print("Training on %d %s" % (device_count, device_type))
+        return self.devices
diff --git a/tango/integrations/flax/train_callback.py b/tango/integrations/flax/train_callback.py
new file mode 100644
index 000000000..2caf9854d
--- /dev/null
+++ b/tango/integrations/flax/train_callback.py
@@ -0,0 +1,194 @@
+import logging
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from tango.common.dataset_dict import DatasetDictBase
+from tango.common.registrable import Registrable
+from tango.workspace import Workspace
+
+from .data import DataLoader
+from .model import Model
+from .optim import Optimizer
+from .train_config import TrainConfig
+
+
+class TrainCallback(Registrable):
+    """
+    A :class:`TrainCallback` is a :class:`~tango.common.Registrable` class
+    that can be used within :class:`FlaxTrainStep` to customize behavior in the training
+    loop. You can set the training callbacks with the ``callbacks`` parameter to :class:`FlaxTrainStep`.
+
+    .. tip::
+        All of the parameters to this base class will be automatically set within
+        the training loop, so you shouldn't include them in your config for your callbacks.
+
+    .. tip::
+        You can access the model being trained through :attr:`self.model <model>`.
+
+    .. important::
+        The ``step`` argument to callback methods is the total/overall number of training steps
+        so far, independent of the current epoch.
+
+    .. seealso::
+        See :class:`~tango.integrations.wandb.WandbTrainCallback` for an example
+        implementation.
+
+    :ivar Workspace workspace: The tango workspace being used.
+    :ivar TrainConfig train_config: The training config.
+    :ivar tango.common.DatasetDictBase dataset_dict: The dataset dict containing train and
+        optional validation splits.
+    :ivar DataLoader train_dataloader: The dataloader used for the training split.
+    :ivar FlaxModel model: The flax model being trained.
+    :ivar Optimizer optimizer: The optimizer being used for training.
+    :ivar DataLoader validation_dataloader: Optional dataloader used for the validation split.
+    """
+
+    def __init__(
+        self,
+        workspace: Workspace,
+        train_config: TrainConfig,
+        dataset: DatasetDictBase,
+        train_dataloader: DataLoader,
+        model: Model,
+        optimizer: Optimizer,
+        validation_dataloader: Optional[DataLoader] = None,
+    ) -> None:
+        self.workspace = workspace
+        self.train_config = train_config
+        self.dataset = dataset
+        self.train_dataloader = train_dataloader
+        self.model = model
+        self.optimizer = optimizer
+        self.validation_dataloader = validation_dataloader
+        self.logger = logging.getLogger(self.__class__.__name__)
+
+    @property
+    def step_id(self) -> str:
+        """
+        The unique ID of the current :class:`~tango.Step`.
+        """
+        return self.train_config.step_id
+
+    @property
+    def step_name(self) -> Optional[str]:
+        """
+        The name of the current:class:`~tango.Step`.
+        """
+        return self.train_config.step_name
+
+    @property
+    def work_dir(self) -> Path:
+        """
+        The working directory of the current train step
+        """
+        return self.train_config.work_dir
+
+    def state_dict(self) -> Dict[str, Any]:
+        """
+        Return any state that needs to be kept after a restart.
+
+        Some callbacks need to maintain state across restarts. This is the callback's opportunity to
+        save it's state. It will be restored using :meth:`load_state_dict`.
+        """
+        return {}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]):
+        """
+        Load the state on a restart.
+
+        Some callbacks need to maintain state across restarts. This is the callback's opportunity to
+        restore it's state. It gets saved using :meth:`state_dict`.
+        """
+        pass
+
+    def pre_train_loop(self) -> None:
+        """
+        Called right before the first batch is processed, or after a restart
+        """
+        pass
+
+    def post_train_loop(self, step: int, epoch: int) -> None:
+        """
+        Called after the training loop completes.
+
+        This is the last method that is called, so any cleanup can be done in this method.
+        """
+        pass
+
+    def pre_epoch(self, step: int, epoch: int) -> None:
+        """
+        Called before start of an epoch. Epochs start at 0.
+        """
+        pass
+
+    def post_epoch(self, step: int, epoch: int) -> None:
+        """
+        Called after an epoch is completed. Epochs start at 0.
+        """
+        pass
+
+    def pre_batch(self, step: int, epoch: int, batch) -> None:
+        """
+        Called directly before processing a batch.
+        """
+
+    def post_batch(self, step: int, epoch: int, train_metrics: Dict) -> None:
+        """
+        Called directly after processing a batch, but before unscaling gradients,
+        clipping gradients, and taking an optimizer step.
+
+        .. note::
+            The ``train_metrics`` here is the dictionary with train metrics of the
+            current batch. If doing, distributed training, use `jax_utils.unreplicate(train_metrics)`
+            before using train_metrics.
+
+            If you need the average loss, use :meth:`log_batch()`.
+        """
+        pass
+
+    def log_batch(self, step: int, epoch: int, train_metrics: Dict) -> None:
+        """
+        Called after the optimizer step. Here ``train_metrics`` is the average metrics across
+        all distributed workers. If doing, distributed training, use
+        `jax_utils.unreplicate(train_metrics)` before using train_metrics.
+
+        .. note::
+            This callback method is not necessarily called on every step.
+            The frequency depends on the value of the ``log_every`` parameter of
+            :class:`FlaxTrainStep`.
+
+        """
+        pass
+
+    def pre_val_loop(self, step: int, val_step: int, state) -> None:
+        """
+        Called right before the validation loop starts.
+        """
+        pass
+
+    def pre_val_batch(self, step: int, val_step: int, epoch: int, val_batch) -> None:
+        """
+        Called right before a validation batch is processed.
+        """
+        pass
+
+    def post_val_batch(self, step: int, val_step: int, epoch: int, val_metrics: Dict) -> None:
+        """
+        Called right after a validation batch is processed with the outputs of the batch.
+
+        .. tip::
+            This method can be used to modify ``val_metrics`` in place, which is useful
+            in scenarios like distributed training where you might need to aggregate metrics
+            in a special way other than a simple average. If that's the case, make sure
+            to set ``auto_aggregate_val_metric`` to ``False`` in :class:`FlaxTrainStep`.
+
+        """
+        pass
+
+    def post_val_loop(
+        self, step: int, epoch: int, val_metric: Optional[float], best_val_metric: Optional[float]
+    ) -> None:
+        """
+        Called right after the evaluation loop finishes
+        """
+        pass
diff --git a/tango/integrations/flax/train_config.py b/tango/integrations/flax/train_config.py
new file mode 100644
index 000000000..8b56e7f25
--- /dev/null
+++ b/tango/integrations/flax/train_config.py
@@ -0,0 +1,128 @@
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class TrainConfig:
+    """
+    Encapsulates the parameters of :class:`FlaxTrainStep`. This is used to pass all the training
+    options to :class:`TrainCallback`.
+    """
+
+    step_id: str
+    """
+    The unique ID of the current step.
+    """
+
+    work_dir: Path
+    """
+    The working directory for the training run.
+    """
+
+    step_name: Optional[str] = None
+    """
+    The name of the current step.
+    """
+
+    train_split: Optional[str] = "train"
+    """
+    The name of the training split.
+    """
+
+    validation_split: Optional[str] = None
+    """
+    The name of the validation split.
+    """
+
+    seed: int = 42
+    """
+    The random seed used to generate
+    """
+
+    train_steps: Optional[int] = None
+    """
+    The number of steps to train for.
+    """
+
+    train_epochs: Optional[int] = None
+    """
+    The number of epochs to train for.
+
+    You cannot specify `train_steps` and `train_epochs` at the same time.
+    """
+
+    validation_steps: Optional[int] = None
+    """
+    The number of validation steps.
+    """
+
+    log_every: int = 10
+    """
+    Controls the frequency of log updates.
+    """
+
+    checkpoint_every: int = 100
+    """
+    Controls the frequency of checkpoints.
+    """
+
+    validate_every: Optional[int] = None
+    """
+    Controls the frequency of the validation loop.
+    """
+
+    is_distributed: bool = False
+    """
+    Whether or not the training job is distributed.
+    """
+
+    val_metric_name: str = "loss"
+    """
+    The name of the validation metric to track.
+    """
+
+    minimize_val_metric: bool = True
+    """
+    Should be ``True`` when the validation metric being tracked should be minimized.
+    """
+
+    auto_aggregate_val_metric: bool = True
+    """
+    Controls automatic aggregation of validation metric.
+    """
+
+    remove_stale_checkpoints: bool = True
+    """
+    Controls removal of stale checkpoints.
+    """
+
+    @property
+    def state_path(self) -> Path:
+        """
+        The path to the latest state checkpoint file.
+        """
+        return self.work_dir / "checkpoint_state_latest"
+
+    @property
+    def best_state_path(self) -> Path:
+        """
+        The path to the best state checkpoint file according to the validation metric or training
+        loss (if no validation split is given).
+        """
+        return self.work_dir / "checkpoint_state_best"
+
+    def should_log_this_step(self, step: int) -> bool:
+        assert self.train_steps is not None
+        return step == 0 or (step + 1) % self.log_every == 0 or step == self.train_steps - 1
+
+    def should_checkpoint_this_step(self, step: int) -> bool:
+        assert self.train_steps is not None
+        return ((step + 1) % self.checkpoint_every == 0) or step == self.train_steps - 1
+
+    def should_log_this_val_step(self, val_step: int) -> bool:
+        assert self.validation_steps is not None
+        return val_step % self.log_every == 0 or val_step == self.validation_steps - 1
+
+    def as_dict(self) -> Dict[str, Any]:
+        return {k: v for k, v in asdict(self).items() if not k.startswith("_")}
diff --git a/tango/integrations/flax/util.py b/tango/integrations/flax/util.py
new file mode 100644
index 000000000..387974e93
--- /dev/null
+++ b/tango/integrations/flax/util.py
@@ -0,0 +1,19 @@
+from typing import Any, Union
+
+import jax
+
+
+def get_PRNGkey(seed: int = 42) -> Union[Any, jax.random.PRNGKeyArray]:
+    """
+    Utility function to create a pseudo-random number generator key
+    given a seed.
+    """
+    return jax.random.PRNGKey(seed)
+
+
+def get_multiple_keys(key, multiple: int = 1) -> Union[Any, jax.random.PRNGKeyArray]:
+    """
+    Utility function to split a PRNG key into multiple new keys.
+    Used in distributed training.
+    """
+    return jax.random.split(key, multiple)
diff --git a/tango/integrations/flax/wrapper.py b/tango/integrations/flax/wrapper.py
new file mode 100644
index 000000000..7ef1e8eae
--- /dev/null
+++ b/tango/integrations/flax/wrapper.py
@@ -0,0 +1,42 @@
+from abc import abstractmethod
+from typing import Dict
+
+from tango.common.registrable import Registrable
+
+
+class FlaxWrapper(Registrable):
+    """
+    A wrapper class which contains functions that need to be defined by the user
+    for using the ``flax::train`` and ``flax::eval`` steps.
+
+    """
+
+    @abstractmethod
+    def train_metrics(self, state, batch, labels) -> Dict:
+        """
+        Returns the train metrics other than loss as Dict.
+        """
+        pass
+
+    @abstractmethod
+    def train_loss(self, params, state, batch, dropout_rng, labels):
+        """
+        This function performs the forward pass and computes loss. The function
+        should return the loss for the batch as a jax device array. The gradient
+        of this function is used for training.
+        """
+        pass
+
+    @abstractmethod
+    def val_metrics(self, batch, logits, labels) -> Dict:
+        """
+        Returns the validation metrics as Dict.
+        """
+        pass
+
+    @abstractmethod
+    def eval_metrics(self, batch, logits, labels) -> Dict:
+        """
+        Returns the evaluation metrics as  Dict.
+        """
+        pass
diff --git a/tango/integrations/transformers/model.py b/tango/integrations/transformers/model.py
index 6be1edb18..475f6aec9 100644
--- a/tango/integrations/transformers/model.py
+++ b/tango/integrations/transformers/model.py
@@ -2,6 +2,13 @@
 
 from transformers.models.auto import modeling_auto
 
+try:
+    from transformers.models.auto import modeling_flax_auto
+
+    from tango.integrations.flax.model import Model as FlaxModel
+except ModuleNotFoundError:
+    pass
+
 from tango.integrations.torch.model import Model
 
 from .config import Config
@@ -31,3 +38,34 @@ def from_config(cls, config: Config, **kwargs) -> Model:
         Model.register("transformers::" + name + "::from_config", constructor="from_config")(
             wrapped_cls
         )
+
+try:
+
+    def flax_auto_model_wrapper_factory(cls: type) -> Type[FlaxModel]:
+        class AutoModelWrapper(cls, FlaxModel):  # type: ignore
+            @classmethod
+            def from_pretrained(
+                cls, pretrained_model_name_or_path: str, config: Optional[Config] = None, **kwargs
+            ) -> FlaxModel:
+                return super().from_pretrained(
+                    pretrained_model_name_or_path, config=config, **kwargs
+                )
+
+            @classmethod
+            def from_config(cls, config: Config, **kwargs) -> FlaxModel:
+                return super().from_config(config, **kwargs)
+
+        return AutoModelWrapper
+
+    for name, cls in modeling_flax_auto.__dict__.items():
+        if isinstance(cls, type) and name.startswith("FlaxAutoModel"):
+            wrapped_cls_ = flax_auto_model_wrapper_factory(cls)
+            FlaxModel.register(
+                "transformers::" + name + "::from_pretrained", constructor="from_pretrained"
+            )(wrapped_cls_)
+            FlaxModel.register(
+                "transformers::" + name + "::from_config", constructor="from_config"
+            )(wrapped_cls_)
+
+except NameError:
+    pass
diff --git a/tango/integrations/wandb/__init__.py b/tango/integrations/wandb/__init__.py
index e56639e55..7851e9435 100644
--- a/tango/integrations/wandb/__init__.py
+++ b/tango/integrations/wandb/__init__.py
@@ -34,10 +34,14 @@
 from .workspace import WandbWorkspace
 
 try:
+    import flax
+    import jax
+    import tensorflow  # flax has a tensorflow dependency
     import torch
 except ModuleNotFoundError:
     pass
 else:
+    from .flax_train_callback import WandbFlaxTrainCallback
     from .torch_train_callback import WandbTrainCallback
 
     __all__.append("WandbTrainCallback")
diff --git a/tango/integrations/wandb/flax_train_callback.py b/tango/integrations/wandb/flax_train_callback.py
new file mode 100644
index 000000000..5c6434898
--- /dev/null
+++ b/tango/integrations/wandb/flax_train_callback.py
@@ -0,0 +1,176 @@
+from typing import Any, Dict, List, Optional
+
+import jax
+import wandb
+from flax import jax_utils
+
+from tango.common.exceptions import ConfigurationError
+from tango.integrations.flax.train_callback import TrainCallback
+
+from .workspace import WandbWorkspace
+
+
+@TrainCallback.register("wandb::log_flax")
+class WandbFlaxTrainCallback(TrainCallback):
+    """
+    A flax :class:`~tango.integrations.flax.TrainCallback` for use with
+    the :class:`~tango.integrations.flax.FlaxTrainStep` that logs training and
+    validation metrics to W&B.
+
+    This can be used with any :class:`~tango.workspace.Workspace` implementation,
+    including :class:`WandbWorkspace`.
+
+    .. tip::
+
+        Registered as a :class:`~tango.integrations.flax.TrainCallback`
+        under the name "wandb::log_flax".
+
+    .. important::
+
+        When this callback is used with the :class:`WandbWorkspace` it will log metrics
+        to the same W&B project that the workspace uses. The ``group`` and ``name``
+        parameters will also automatically be set, so a :class:`~tango.common.exceptions.ConfigurationError`
+        will be raised if any of ``project``, ``entity``, ``group``, or ``name`` are set in this callback.
+
+    :param project:
+            W&B project to associated this run with.
+
+    :param entity:
+        W&B entity (user or organization) to associated this run with.
+
+    :param group:
+        W&B group to associated this run with.
+
+    :param name:
+        Set the name of the run in W&B. If not set, the default will be the name of the step.
+
+    :param notes:
+        Arbitrary notes to add in W&B to this run.
+
+    :param tags:
+        Arbitrary tags to add in W&B to this run.
+
+    :param watch_model:
+        If ``True``, ``wandb.watch()`` is called to collect gradients and other information
+        about the model throughout training.
+        See `docs.wandb.ai/ref/python/watch <https://docs.wandb.ai/ref/python/watch>`_.
+
+    :param wandb_config:
+        Arbitrary configuration fields to set in W&B for this run.
+        See `docs.wandb.ai/guides/track/config <https://docs.wandb.ai/guides/track/config>`_.
+
+    """
+
+    def __init__(
+        self,
+        *args,
+        project: Optional[str] = None,
+        entity: Optional[str] = None,
+        group: Optional[str] = None,
+        name: Optional[str] = None,
+        notes: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        watch_model: bool = False,
+        wandb_config: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+
+        if isinstance(self.workspace, WandbWorkspace) or wandb.run is not None:
+            err_msg_template = "Cannot set '{var_name}' in WandbTrainCallback "
+            if isinstance(self.workspace, WandbWorkspace):
+                err_msg_template += "since it has already been set from the WandbWorkspace."
+            else:
+                err_msg_template += "since a W&B run has already been initialized."
+            for var, var_name in [
+                (project, "project"),
+                (entity, "entity"),
+                (group, "group"),
+                (name, "name"),
+            ]:
+                if var is not None:
+                    raise ConfigurationError(err_msg_template.format(var_name=var_name))
+
+        self.project = (
+            project if not isinstance(self.workspace, WandbWorkspace) else self.workspace.project
+        )
+        self.entity = (
+            entity if not isinstance(self.workspace, WandbWorkspace) else self.workspace.entity
+        )
+        self.group = group or self.step_id
+        self.notes = notes
+        self.tags = tags
+        self.watch_model = watch_model
+        self.wandb_config = self.train_config.as_dict()
+
+        if wandb_config is not None:
+            self.wandb_config.update(wandb_config)
+        if wandb.run is None:
+            self.wandb_config["job_type"] = "train_metrics"
+
+        self.run_name: str = name or self.step_name or "train"
+
+        self.run_id: str = (
+            wandb.run.id if wandb.run is not None else self.step_id  # type: ignore[attr-defined]
+        )
+        self.resume: Optional[str] = None
+        self.should_finalize_run: bool = (
+            wandb.run is None
+        )  # if we have to start out own W&B run, we need to finish it
+
+    def state_dict(self) -> Dict[str, Any]:
+        return {}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        self.resume = "allow"
+
+    def pre_train_loop(self) -> None:
+        if wandb.run is None:
+            if self.run_id is None:
+                self.run_id = self.step_id
+
+            wandb.init(
+                id=self.run_id,
+                dir=str(self.work_dir),
+                project=self.project,
+                entity=self.entity,
+                group=self.group,
+                name=self.run_name,
+                notes=self.notes,
+                config=self.wandb_config,
+                tags=self.tags,
+                job_type="train_metrics",
+            )
+        else:
+            # We are already running inside of a W&B run, possibly because
+            # we're using the WandbWorkspace.
+            wandb.config.update(self.wandb_config)
+            if self.tags:
+                wandb.run.tags = (wandb.run.tags or tuple()) + tuple(self.tags)
+            if self.notes:
+                wandb.run.notes = self.notes
+
+        if self.watch_model:
+            wandb.watch(self.model)
+
+    def post_train_loop(self, step: int, epoch: int) -> None:
+        if self.should_finalize_run:
+            wandb.finish()
+
+    def log_batch(self, step: int, epoch: int, train_metrics: Dict) -> None:
+        if len(jax.devices()) > 1:
+            train_metrics = jax_utils.unreplicate(train_metrics)
+        metrics = {"train/loss": train_metrics["loss"], "epoch": epoch}
+        wandb.log(metrics, step=step + 1)
+
+    def post_val_loop(
+        self, step: int, epoch: int, val_metric: Optional[float], best_val_metric: Optional[float]
+    ) -> None:
+        wandb.log(
+            {
+                f"val/{self.train_config.val_metric_name}": val_metric,
+                f"val/best_{self.train_config.val_metric_name}": best_val_metric,
+                "epoch": epoch,
+            },
+            step=step + 1,
+        )


Test patch:
diff --git a/test_fixtures/integrations/flax/__init__.py b/test_fixtures/integrations/flax/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/test_fixtures/integrations/flax/config.jsonnet b/test_fixtures/integrations/flax/config.jsonnet
new file mode 100644
index 000000000..90180158e
--- /dev/null
+++ b/test_fixtures/integrations/flax/config.jsonnet
@@ -0,0 +1,52 @@
+{
+    "steps": {
+        "data_full": {
+            "type": "datasets::load",
+            "path": "iohadrubin/mini_xsum",
+        },
+        "data": {
+            "type": "datasets::dataset_remix",
+            "input": {"type": "ref", "ref": "data_full"},
+            "new_splits": {"train": "train[:20]", "validation": "validation[:20]"},
+        },
+        "tokenize": {
+            "type": "tokenize_data",
+            "dataset": {
+                "type": "ref",
+                "ref": "data"
+            }
+        },
+        "train": {
+            "type": "flax::train",
+            "model": {
+                "type" : "transformers::FlaxAutoModelForSeq2SeqLM::from_pretrained",
+                "pretrained_model_name_or_path" : "t5-small"
+            },
+            "dataset": {
+                "type": "ref",
+                "ref": "tokenize"
+            },
+            "optimizer": {
+                "type" : "optax::adamw",
+                "learning_rate" : 2e-5
+            },
+            "train_dataloader": {
+                "batch_size": 16,
+                "drop_last": true
+            },
+            "wrapper": {
+                "type": "xsum_wrapper"
+            },
+            "train_split": "train",
+            "validation_split" : "validation",
+            "validate_every" : 1,
+            "validation_dataloader": {
+                "batch_size": 16,
+                "drop_last": true
+            },
+            "train_epoch": 1,
+            "checkpoint_every": 1,
+            "log_every": 1
+        }
+    }
+}
\ No newline at end of file
diff --git a/test_fixtures/integrations/flax/xsum.py b/test_fixtures/integrations/flax/xsum.py
new file mode 100644
index 000000000..5d2f55f96
--- /dev/null
+++ b/test_fixtures/integrations/flax/xsum.py
@@ -0,0 +1,113 @@
+import jax.numpy as jnp
+import numpy as np
+import optax
+from flax.training.common_utils import onehot
+from transformers import AutoConfig, AutoTokenizer, FlaxAutoModelForSeq2SeqLM
+
+from tango.integrations.flax import FlaxWrapper
+from tango.step import Step
+
+"""
+A minimal xsum t5-small config for testing.
+"""
+
+
+@Step.register("tokenize_data")
+class PreProcessing(Step):
+    DETERMINISTIC = False
+
+    def run(self, dataset):
+        tokenizer = AutoTokenizer.from_pretrained("t5-small")
+        model = FlaxAutoModelForSeq2SeqLM.from_pretrained("t5-small")
+        model_module = __import__(model.__module__, fromlist=["shift_tokens_tight"])
+        shift_tokens_right_fn = getattr(model_module, "shift_tokens_right")
+        config = AutoConfig.from_pretrained("t5-small")
+
+        MAX_SOURCE_LENGTH = 512
+        MAX_TGT_LENGTH = 64
+
+        def preprocess_function(examples):
+            inputs = examples["document"]
+            targets = examples["summary"]
+            inputs = [inp for inp in inputs]
+            model_inputs = tokenizer(
+                inputs,
+                max_length=MAX_SOURCE_LENGTH,
+                padding="max_length",
+                truncation=True,
+                return_tensors="np",
+            )
+
+            # Setup the tokenizer for targets
+            with tokenizer.as_target_tokenizer():
+                labels = tokenizer(
+                    targets,
+                    max_length=MAX_TGT_LENGTH,
+                    padding="max_length",
+                    truncation=True,
+                    return_tensors="np",
+                )
+
+            model_inputs["labels"] = labels["input_ids"]
+            decoder_input_ids = shift_tokens_right_fn(
+                labels["input_ids"], config.pad_token_id, config.decoder_start_token_id
+            )
+            model_inputs["decoder_input_ids"] = np.asarray(decoder_input_ids)
+
+            # We need decoder_attention_mask so we can ignore pad tokens from loss
+            model_inputs["decoder_attention_mask"] = labels["attention_mask"]
+
+            return model_inputs
+
+        column_names = dataset["train"].column_names
+
+        dataset = dataset.map(
+            preprocess_function,
+            batched=True,
+            remove_columns=column_names,
+            desc="Running tokenizer on dataset",
+        )
+
+        return dataset
+
+
+@FlaxWrapper.register("xsum_wrapper")  # type: ignore
+class TransformerWrapper(FlaxWrapper):
+    def compute_metrics(self, state, batch, labels):
+        # return empty dict if no other metrics to compute
+        return {}
+
+    def loss_helper(self, logits, labels, batch):
+        label_smoothing_factor = 0
+        padding_mask = batch["decoder_attention_mask"]
+        vocab_size = logits.shape[-1]
+        confidence = 1.0 - label_smoothing_factor
+        low_confidence = (1.0 - confidence) / (vocab_size - 1)
+        normalizing_constant = -(
+            confidence * jnp.log(confidence)
+            + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)
+        )
+        soft_labels = onehot(labels, vocab_size, on_value=confidence, off_value=low_confidence)
+
+        loss = optax.softmax_cross_entropy(logits, soft_labels)
+        loss = loss - normalizing_constant
+
+        # ignore padded tokens from loss
+        loss = loss * padding_mask
+        loss = loss.sum() / padding_mask.sum()
+        return loss
+
+    def train_loss(self, params, state, batch, dropout_rng, labels):
+        logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
+        loss = self.loss_helper(logits, labels, batch)
+        return loss
+
+    def val_metrics(self, batch, logits, labels):
+        loss = self.loss_helper(logits, labels, batch)
+        metrics = {"loss": loss}
+        return metrics
+
+    def eval_metrics(self, batch, logits, labels):
+        loss = self.loss_helper(logits, labels, batch)
+        metrics = {"loss": loss}
+        return metrics
diff --git a/tests/integrations/flax/__init__.py b/tests/integrations/flax/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/integrations/flax/data_test.py b/tests/integrations/flax/data_test.py
new file mode 100644
index 000000000..930c41518
--- /dev/null
+++ b/tests/integrations/flax/data_test.py
@@ -0,0 +1,35 @@
+from typing import Dict
+
+from transformers import AutoTokenizer
+
+from tango.common.testing import TangoTestCase
+from tango.integrations.flax import DataLoader, FlaxDataLoader
+from tango.integrations.flax.util import get_PRNGkey
+from tango.step import Step
+
+
+class TestDataStep(TangoTestCase):
+    def test_dataloader(self) -> None:
+        assert "flax::dataloader" in DataLoader.list_available()
+
+    def test_sample_data(self) -> None:
+        step = Step.from_params(  # type: ignore[assignment]
+            {
+                "type": "datasets::load",
+                "path": "lhoestq/demo1",
+                "split": "train",
+                "cache_dir": str(self.TEST_DIR / "cache"),
+            }
+        )
+
+        dataset = step.result()
+        tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+        column_names = dataset.column_names
+        dataset = dataset.map(
+            lambda e: tokenizer(e["review"], truncation=True, padding="max_length")
+        )
+        dataset = dataset.remove_columns(column_names)
+        data = FlaxDataLoader(dataset, batch_size=16)
+        rng = get_PRNGkey()
+        for batch in data(rng, do_distributed=False):
+            assert isinstance(batch, Dict)
diff --git a/tests/integrations/flax/format_test.py b/tests/integrations/flax/format_test.py
new file mode 100644
index 000000000..b2b1742f9
--- /dev/null
+++ b/tests/integrations/flax/format_test.py
@@ -0,0 +1,14 @@
+import os
+
+from tango import Format
+from tango.common.testing import TangoTestCase
+from tango.integrations.flax.format import FlaxFormat
+
+
+class TestTorchFormat(TangoTestCase):
+    def test_read_write(self):
+        flax_format: FlaxFormat = Format.by_name("flax")()  # type: ignore[assignment]
+        flax_format.write({"a": 1}, self.TEST_DIR)
+        assert os.path.exists(self.TEST_DIR / "checkpoint_0")
+        data = flax_format.read(self.TEST_DIR)
+        assert data == {"a": 1}
diff --git a/tests/integrations/flax/optim_test.py b/tests/integrations/flax/optim_test.py
new file mode 100644
index 000000000..54028d1b4
--- /dev/null
+++ b/tests/integrations/flax/optim_test.py
@@ -0,0 +1,9 @@
+from tango.integrations.flax.optim import LRScheduler, Optimizer
+
+
+def test_all_optimizers_registered():
+    assert "optax::adafactor" in Optimizer.list_available()
+
+
+def test_all_lr_schedulers_registered():
+    assert "optax::constant_schedule" in LRScheduler.list_available()
diff --git a/tests/integrations/flax/train_test.py b/tests/integrations/flax/train_test.py
new file mode 100644
index 000000000..0eb6bcebf
--- /dev/null
+++ b/tests/integrations/flax/train_test.py
@@ -0,0 +1,24 @@
+from tango.common.logging import initialize_logging, teardown_logging
+from tango.common.testing import TangoTestCase
+
+
+class TestTrainStep(TangoTestCase):
+    def setup_method(self):
+        super().setup_method()
+        initialize_logging(enable_cli_logs=True)
+
+    def teardown_method(self):
+        super().teardown_method()
+        teardown_logging()
+
+    def test_trainer(self):
+        result_dir = self.run(
+            self.FIXTURES_ROOT / "integrations" / "flax" / "config.jsonnet",
+            include_package=[
+                "test_fixtures.integrations.common",
+                "test_fixtures.integrations.flax",
+            ],
+        )
+        assert (
+            result_dir / "train" / "work" / "checkpoint_state_latest" / "checkpoint_0"
+        ).is_file()