LICENSE
README.md
pyproject.toml
jax_tpu_embedding/__init__.py
jax_tpu_embedding.egg-info/PKG-INFO
jax_tpu_embedding.egg-info/SOURCES.txt
jax_tpu_embedding.egg-info/dependency_links.txt
jax_tpu_embedding.egg-info/requires.txt
jax_tpu_embedding.egg-info/top_level.txt
jax_tpu_embedding/sparsecore/__init__.py
jax_tpu_embedding/sparsecore/version.py
jax_tpu_embedding/sparsecore/lib/__init__.py
jax_tpu_embedding/sparsecore/lib/core/__init__.py
jax_tpu_embedding/sparsecore/lib/core/constants.py
jax_tpu_embedding/sparsecore/lib/core/input_preprocessing.py
jax_tpu_embedding/sparsecore/lib/core/libinput_preprocessing_core.so
jax_tpu_embedding/sparsecore/lib/core/libinput_preprocessing_threads.so
jax_tpu_embedding/sparsecore/lib/core/libinput_preprocessing_util.so
jax_tpu_embedding/sparsecore/lib/core/pybind_input_preprocessing.so
jax_tpu_embedding/sparsecore/lib/core/primitives/__init__.py
jax_tpu_embedding/sparsecore/lib/core/primitives/optimizers_computation.py
jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_csr.py
jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_adagrad.py
jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_adagrad_momentum.py
jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_adam.py
jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_ftrl.py
jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_laprop.py
jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_sgd.py
jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_optimizer_grad.py
jax_tpu_embedding/sparsecore/lib/core/primitives/utils.py
jax_tpu_embedding/sparsecore/lib/fdo/__init__.py
jax_tpu_embedding/sparsecore/lib/fdo/fdo_client.py
jax_tpu_embedding/sparsecore/lib/fdo/fdo_utils.py
jax_tpu_embedding/sparsecore/lib/fdo/file_fdo_client.py
jax_tpu_embedding/sparsecore/lib/flax/__init__.py
jax_tpu_embedding/sparsecore/lib/flax/embed.py
jax_tpu_embedding/sparsecore/lib/flax/embed_optimizer.py
jax_tpu_embedding/sparsecore/lib/nn/__init__.py
jax_tpu_embedding/sparsecore/lib/nn/embedding.py
jax_tpu_embedding/sparsecore/lib/nn/embedding_spec.py
jax_tpu_embedding/sparsecore/lib/nn/table_stacking.py
jax_tpu_embedding/sparsecore/lib/proto/__init__.py
jax_tpu_embedding/sparsecore/lib/proto/embedding_spec_pb2.py
jax_tpu_embedding/sparsecore/utils/__init__.py
jax_tpu_embedding/sparsecore/utils/utils.py