jax-tpu-embedding 0.1.0__tar.gz

jax_tpu_embedding-0.1.0/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

jax_tpu_embedding-0.1.0/PKG-INFO

@@ -0,0 +1,71 @@
+Metadata-Version: 2.4
+Name: jax_tpu_embedding
+Version: 0.1.0
+Summary: JAX SparseCore API
+Author-email: JAX SC Authors <jax-sc-dev@google.com>
+License-Expression: Apache-2.0
+Project-URL: homepage, https://github.com/jax-ml/jax-tpu-embedding
+Project-URL: repository, https://github.com/jax-ml/jax-tpu-embedding
+Project-URL: issue_tracker, https://github.com/jax-ml/jax-tpu-embedding/issues
+Keywords: jax,tpu,embedding,sparsecore,machine learning,sparse tensor,distributed computing
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: absl-py
+Requires-Dist: einops
+Requires-Dist: flax
+Requires-Dist: jax
+Requires-Dist: numpy
+Requires-Dist: portpicker
+Dynamic: license-file
+
+# JAX TPU Embedding (JAX SparseCore)
+
+[![Build and test](https://github.com/jax-ml/jax-tpu-embedding/actions/workflows/build_and_test.yml/badge.svg)](https://github.com/jax-ml/jax-tpu-embedding/actions/workflows/build_and_test.yml)
+[![PyPI version](https://img.shields.io/pypi/v/jax-tpu-embedding)](https://pypi.org/project/jax-tpu-embedding/)
+[![Documentation Status](https://readthedocs.org/projects/jax-tpu-embedding/badge/?version=latest)](https://jax-tpu-embedding.readthedocs.io/en/latest/?badge=latest)
+
+[**Quick Overview**](https://jax-tpu-embedding.readthedocs.io/en/latest/overview.html)
+| [**Install guide**](#installation)
+| [**Reference docs**](https://jax-tpu-embedding.readthedocs.io/en/latest/)
+
+## What is JAX TPU Embedding?
+
+JAX SparseCore provides support for leveraging the SparseCore accelerators
+present in TPU generations starting with TPU v5.
+
+SparseCores are specialized tiled processors engineered for high-performance
+acceleration of workloads that involve irregular, sparse memory access and
+computation. While they excel at tasks like embedding lookups (common in deep
+learning recommendation models), their capabilities extend to accelerating a
+variety of other dynamic and sparse workloads on large datasets stored in High
+Bandwidth Memory (HBM).
+
+This is a research project, not an official Google product. Expect sharp edges.
+Please help by trying it out, reporting bugs, and letting us know what you
+think!
+
+## Installation
+
+You can install JAX TPU Embedding from PyPI:
+
+```bash
+pip install jax-tpu-embedding
+```
+
+*Note: To use TPU acceleration, you must run in an environment with access to TPU v5+ hardware and have the appropriate `jax` and `jaxlib` TPU releases installed (see the [JAX installation guide](https://github.com/google/jax#installation)).*
+
+## Documentation
+
+For detailed guides, specifications, and tutorials, see the [JAX TPU Embedding Documentation](https://jax-tpu-embedding.readthedocs.io/en/latest/).
+
+*   [Quick Overview for Users](https://jax-tpu-embedding.readthedocs.io/en/latest/overview.html)
+*   [Embedding Specification](https://jax-tpu-embedding.readthedocs.io/en/latest/embedding.html)
+*   [Input Processing](https://jax-tpu-embedding.readthedocs.io/en/latest/input_processing.html)
+*   [Tutorials (Flax & Primitives)](https://jax-tpu-embedding.readthedocs.io/en/latest/#tutorials)

jax_tpu_embedding-0.1.0/README.md

@@ -0,0 +1,44 @@
+# JAX TPU Embedding (JAX SparseCore)
+
+[![Build and test](https://github.com/jax-ml/jax-tpu-embedding/actions/workflows/build_and_test.yml/badge.svg)](https://github.com/jax-ml/jax-tpu-embedding/actions/workflows/build_and_test.yml)
+[![PyPI version](https://img.shields.io/pypi/v/jax-tpu-embedding)](https://pypi.org/project/jax-tpu-embedding/)
+[![Documentation Status](https://readthedocs.org/projects/jax-tpu-embedding/badge/?version=latest)](https://jax-tpu-embedding.readthedocs.io/en/latest/?badge=latest)
+
+[**Quick Overview**](https://jax-tpu-embedding.readthedocs.io/en/latest/overview.html)
+| [**Install guide**](#installation)
+| [**Reference docs**](https://jax-tpu-embedding.readthedocs.io/en/latest/)
+
+## What is JAX TPU Embedding?
+
+JAX SparseCore provides support for leveraging the SparseCore accelerators
+present in TPU generations starting with TPU v5.
+
+SparseCores are specialized tiled processors engineered for high-performance
+acceleration of workloads that involve irregular, sparse memory access and
+computation. While they excel at tasks like embedding lookups (common in deep
+learning recommendation models), their capabilities extend to accelerating a
+variety of other dynamic and sparse workloads on large datasets stored in High
+Bandwidth Memory (HBM).
+
+This is a research project, not an official Google product. Expect sharp edges.
+Please help by trying it out, reporting bugs, and letting us know what you
+think!
+
+## Installation
+
+You can install JAX TPU Embedding from PyPI:
+
+```bash
+pip install jax-tpu-embedding
+```
+
+*Note: To use TPU acceleration, you must run in an environment with access to TPU v5+ hardware and have the appropriate `jax` and `jaxlib` TPU releases installed (see the [JAX installation guide](https://github.com/google/jax#installation)).*
+
+## Documentation
+
+For detailed guides, specifications, and tutorials, see the [JAX TPU Embedding Documentation](https://jax-tpu-embedding.readthedocs.io/en/latest/).
+
+*   [Quick Overview for Users](https://jax-tpu-embedding.readthedocs.io/en/latest/overview.html)
+*   [Embedding Specification](https://jax-tpu-embedding.readthedocs.io/en/latest/embedding.html)
+*   [Input Processing](https://jax-tpu-embedding.readthedocs.io/en/latest/input_processing.html)
+*   [Tutorials (Flax & Primitives)](https://jax-tpu-embedding.readthedocs.io/en/latest/#tutorials)

jax_tpu_embedding-0.1.0/jax_tpu_embedding/sparsecore/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2024 The JAX SC Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""JAX SparseCore library."""
+
+from jax_tpu_embedding.sparsecore import version
+
+__version__ = version.__version__

jax_tpu_embedding-0.1.0/jax_tpu_embedding/sparsecore/lib/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 The JAX SC Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Empty file needed by setuptools.find_packages to recognize this as a package.

jax_tpu_embedding-0.1.0/jax_tpu_embedding/sparsecore/lib/auto_pipelining/auto_pipelining.py

@@ -0,0 +1,273 @@
+# Copyright 2024 The JAX SC Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Autopipelining decorator.
+
+This decorator is used to automatically transform a train step function to use
+SparseCore pipelining.
+
+Example:
+
+```python
+@auto_pipelining.auto_pipelining
+def train_step(carry: Carry, ...) -> tuple[Carry, ...]:
+  ...
+  return new_carry, ...
+```
+
+With the decorator, the `train_step` function will be transformed to take
+pipelining state as the first argument, before the carry. It also returns a new
+pipelining state as the first return value.
+
+```python
+def train_step(state: EmbeddingPipeliningState, carry: Carry, ...) ->
+    tuple[EmbeddingPipeliningState, Carry, ...]:
+  ...
+  return new_pipelining_state, new_carry, ...
+```
+
+**Notice** that we expect that the train step function returns multiple values,
+with carry as the first value. We don't support the train_step that only returns
+the carry. This is because the carry may be defined as a NamedTuple. We can not
+diferentiate a single carry result, vs multiple results.
+
+
+# Carry
+
+The training function should take carry as the first argument, and return a new
+carry as the first return value.
+
+The carry should include all the training states. Specifically, embedding tables
+should be part of carry.
+
+Usually, the carry should be donated to the train step function in each step.
+And the function should return a new updated carry with exactly the same
+structure. There's usually only one carry instance at one time.
+
+On the other hand, args should be the "data" fed into each train step. It
+should be ok for the pipelining system to keep them for up to three steps.
+
+A typical training loop looks like this:
+
+```
+carry = init_carry()
+for (data1, data2) in dataset:
+  carry, loss, metrics = train_step(carry, data1, data2)
+```
+
+With pipelining, the training loop should be modified to pass pipelining state:
+
+```
+carry = init_carry()
+state = auto_pipelining.init_state()
+for (data1, data2) in dataset:
+  state, (carry, loss, metrics) = train_step(state, carry, data1, data2)
+```
+
+This carry assumption may be avoidable, but it greatly simplifies the
+implementation and use of auto-pipelining system:
+
+1. Training data and training state are naturally separated. We know what part
+of the input should be reused across steps, and what part is step specific.
+
+2. For all training state arrays, we can easily map them from inputs to the
+outputs. This helps create a "no-op" training step. This is useful for the first
+and last few steps, where we only run part of the SparseCore forward, dense,
+or the SparseCore backward.
+
+
+# Limitations on embedding lookups and updates
+
+1. There should be only one embedding lookup and embedding update for each
+embedding table.
+
+2. The embedding lookups and updates should not be wrapped in any ops, e.g.
+`jax.lax.cond`, or nested `jax.jit`.
+
+3. Each embedding lookup and update shard_map should not contain other
+operations.
+
+4. Embedding table should be the last argument to lookup and update shard_map.
+
+The easiest way to comply to these limitations is to use our Flax API.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import functools
+import logging
+import os
+from typing import Any, Callable, Concatenate, NamedTuple, ParamSpec, TypeVar
+
+import jax
+import jax.extend as jex
+from jax_tpu_embedding.sparsecore.lib.auto_pipelining import decompose
+
+
+Carry = decompose.Carry
+
+
+def dump_jaxpr(jaxpr: jex.core.ClosedJaxpr | jex.core.Jaxpr, name: str) -> None:
+  """Dumps a Jaxpr to a file, if the TEST_UNDECLARED_OUTPUTS_DIR is set."""
+  dirname = os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR')
+  if not dirname:
+    return
+  dump_path = os.path.join(dirname, f'jaxpr_{name}.txt')
+
+  logging.info('Writing %s Jaxpr to %s', name, dump_path)
+
+  with open(dump_path, 'w') as f:
+    f.write(jaxpr.pretty_print(source_info=True))
+
+
+class EmbeddingPipeliningState(NamedTuple):
+  activations_1: list[jax.Array] | None = None
+  args_1: Any = None
+  updates_0: list[jax.Array] | None = None
+
+
+@dataclasses.dataclass
+class PipeliningFunction:
+  train_step_func: Callable[..., Any]
+  finalize: Callable[..., Any]
+
+  def __call__(self, state: EmbeddingPipeliningState, *args_2):
+    return self.train_step_func(state, *args_2)
+
+
+def init_state() -> EmbeddingPipeliningState:
+  """Initializes the EmbeddingPipeliningState."""
+  return EmbeddingPipeliningState()
+
+
+def _run(
+    runner: decompose.FunctionRunner,
+    state: EmbeddingPipeliningState,
+    carry: Carry,
+    args_2,
+):
+  """Runs one step of the pipelined training loop."""
+  if args_2 is not None:
+    if state.updates_0 is not None:
+      # Run combined update and lookup (SparseCore backward then forward).
+      activations_2, carry = runner.embedding_update_lookup(
+          state.updates_0, carry, *args_2
+      )
+    else:
+      # Run lookup only (SparseCore forward).
+      activations_2 = runner.embedding_lookup(carry, *args_2)
+  else:
+    activations_2 = None
+    if state.updates_0 is not None:
+      # Run update only (SparseCore backward).
+      carry = runner.embedding_update(state.updates_0, carry)
+
+  if state.args_1 is not None and state.activations_1 is not None:
+    # Run dense layer.
+    updates_1, carry, *func_results = runner.dense(
+        state.activations_1, carry, *state.args_1
+    )
+  else:
+    # No dense computation needed, create empty results.
+    updates_1 = None
+    _, *func_results = runner.none_result()
+
+  # Prepare data for the next step.
+  next_state = EmbeddingPipeliningState(
+      activations_1=activations_2,
+      # During finalization, args_2 is None. Keep the last args_1 for tracing
+      # This can be removed if we move to one program setting.
+      args_1=args_2 or state.args_1,
+      updates_0=updates_1,
+  )
+
+  return next_state, (
+      carry,
+      *func_results,
+  )
+
+
+def finalize_pipelining(
+    fn: Callable[..., Any], state: EmbeddingPipeliningState, carry: Carry
+):
+  """Finalizes the pipelining run."""
+  assert isinstance(
+      fn, PipeliningFunction
+  ), 'The function is not auto-pipelining wrapped.'
+  state, (carry, *res) = fn.finalize(state, carry)
+  state, (carry, *_) = fn.finalize(state, carry)
+  return state, (carry, *res)
+
+
+FunctionArgs = ParamSpec('FunctionArgs')
+FunctionResults = TypeVar('FunctionResults')
+
+
+def auto_pipelining(
+    fn: Callable[FunctionArgs, FunctionResults],
+    out_shardings: tuple[Any, ...],
+    static_argnums: tuple[int, ...] = (),
+) -> Callable[
+    Concatenate[EmbeddingPipeliningState, FunctionArgs],
+    tuple[EmbeddingPipeliningState, FunctionResults],
+]:
+  """A decorator to transform a function for SparseCore pipelining."""
+
+  def _build_runner(carry: Carry, args_2):
+    """Compile the function if it's not compiled yet."""
+    closed_jaxpr, res_shape = jax.make_jaxpr(
+        fn, static_argnums=static_argnums, return_shape=True
+    )(carry, *args_2)
+    dump_jaxpr(closed_jaxpr, 'original')
+    assert isinstance(res_shape, tuple), 'train_step should return a tuple.'
+    assert jax.tree.structure(carry) == jax.tree.structure(res_shape[0]), (
+        'The structure of the first argument (carry) should be the same as'
+        ' the first return value.'
+    )
+
+    res_structure = jax.tree.structure(res_shape)
+    carry_structure = jax.tree.structure(carry)
+    runner = decompose.decompose(closed_jaxpr, res_structure, carry_structure)
+
+    dump_jaxpr(runner.lookup_jaxpr, 'lookup')
+    dump_jaxpr(runner.dense_jaxpr, 'dense')
+    dump_jaxpr(runner.update_jaxpr, 'update')
+    dump_jaxpr(runner.update_lookup_jaxpr, 'update_lookup')
+
+    return runner
+
+  @functools.wraps(fn)
+  def train_pipeline(state: EmbeddingPipeliningState, carry: Carry, *args_2):
+    runner = _build_runner(carry, args_2)
+    return _run(runner, state, carry, args_2)
+
+  def finalize_pipeline(state: EmbeddingPipeliningState, carry: Carry):
+    assert state.args_1, (
+        'Args_1 should not be None. This is likely caused by a bug in the'
+        ' auto-pipelining system.'
+    )
+
+    runner = _build_runner(carry, state.args_1)
+    return _run(runner, state, carry, None)
+
+  train_pipeline = jax.jit(
+      train_pipeline,
+      out_shardings=out_shardings,
+      donate_argnames=('carry',),
+  )
+  finalize_pipeline = jax.jit(
+      finalize_pipeline, out_shardings=out_shardings, donate_argnames=('carry',)
+  )
+
+  return PipeliningFunction(train_pipeline, finalize_pipeline)