Introduction

Partitioned variables are a way to split one logical TensorFlow variable into multiple smaller shards. They are common in older large-model and distributed-training setups, especially for embeddings. Saving and restoring them works reliably, but only if the checkpoint naming and variable construction stay consistent between the save side and the restore side.

What a Partitioned Variable Actually Is

In TensorFlow 1 style code, a partitioned variable is usually created through tf.compat.v1.get_variable with a partitioner. TensorFlow then creates several underlying shard variables while exposing one logical grouped object to your graph-building code.

A simple example:

1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5with tf.compat.v1.variable_scope(
6    "demo",
7    partitioner=tf.compat.v1.fixed_size_partitioner(num_shards=2)
8):
9    embedding = tf.compat.v1.get_variable(
10        "embedding",
11        shape=[6, 4],
12        initializer=tf.compat.v1.ones_initializer(),
13    )
14
15print(embedding)

The important point is that the checkpoint will store shards such as embedding/part_0 and embedding/part_1, not one monolithic tensor entry.

Saving in TensorFlow 1 Style

If you are using graph mode and Saver, the usual pattern is straightforward:

1import os
2import tensorflow as tf
3
4tf.compat.v1.disable_eager_execution()
5
6checkpoint_dir = "/tmp/tf_partitioned_demo"
7os.makedirs(checkpoint_dir, exist_ok=True)
8
9with tf.compat.v1.variable_scope(
10    "demo",
11    partitioner=tf.compat.v1.fixed_size_partitioner(num_shards=2)
12):
13    embedding = tf.compat.v1.get_variable(
14        "embedding",
15        shape=[6, 4],
16        initializer=tf.compat.v1.ones_initializer(),
17    )
18
19saver = tf.compat.v1.train.Saver()
20
21with tf.compat.v1.Session() as sess:
22    sess.run(tf.compat.v1.global_variables_initializer())
23    path = saver.save(sess, os.path.join(checkpoint_dir, "model.ckpt"))
24    print(path)

Saver knows how to write the partitioned shards, so you do not need custom save logic for the common case.

Restoring Requires the Same Variable Structure

The restore side must build variables with the same scope names, base variable names, and partitioning scheme. If you change those, TensorFlow no longer knows how the checkpoint shards map back into the graph.

1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5with tf.compat.v1.variable_scope(
6    "demo",
7    partitioner=tf.compat.v1.fixed_size_partitioner(num_shards=2)
8):
9    embedding = tf.compat.v1.get_variable(
10        "embedding",
11        shape=[6, 4],
12        initializer=tf.compat.v1.zeros_initializer(),
13    )
14
15saver = tf.compat.v1.train.Saver()
16
17with tf.compat.v1.Session() as sess:
18    saver.restore(sess, "/tmp/tf_partitioned_demo/model.ckpt")
19    print(sess.run(embedding))

The initializer can differ, because restore overwrites the variable values. The structure cannot differ.

The Most Common Restore Failure

The usual failure is not a broken checkpoint. It is a mismatch such as:

• different variable scope
• different base variable name
• different number of shards
• different shape

For example, saving with 2 shards and restoring with 4 shards is usually not a safe drop-in change. The checkpoint keys no longer line up with the graph in the way Saver expects.

Inspect Checkpoint Contents When Debugging

If restore fails, inspect the checkpoint names directly:

1import tensorflow as tf
2
3vars_in_ckpt = tf.train.list_variables("/tmp/tf_partitioned_demo/model.ckpt")
4for name, shape in vars_in_ckpt:
5    print(name, shape)

This is one of the fastest ways to confirm what TensorFlow actually saved. You can then compare those names against the graph variables you are creating during restore.

TensorFlow 2 Perspective

In TensorFlow 2, the usual recommendation is to use object-based checkpointing through tf.train.Checkpoint or Keras saving APIs. Partitioned-variable workflows are much more tied to TensorFlow 1 style graph code and large older models.

That means if you are maintaining legacy code, tf.compat.v1 compatibility mode is often the practical answer. If you are starting fresh, design around modern checkpointing rather than older partitioned-variable patterns unless you have a strong reason not to.

A Good Rule for Stable Checkpoints

Treat the save and restore graph definitions as part of the checkpoint contract. If you need to refactor:

1. inspect the existing checkpoint names
2. preserve scope and variable naming where possible
3. only change partitioning layout if you are also planning a controlled checkpoint conversion

That mindset prevents a lot of avoidable restore failures.

Common Pitfalls

• Changing the variable scope or variable name between save and restore.
• Rebuilding the model with a different partitioning scheme and expecting the checkpoint to restore automatically.
• Debugging restore failures without listing checkpoint variable names first.
• Mixing TensorFlow 1 style Saver assumptions with TensorFlow 2 object-based checkpoint expectations.
• Assuming the logical grouped variable name is the only name that matters, while the checkpoint stores shard names too.

Summary

• Partitioned variables are saved as multiple shards under one logical variable concept.
• 'tf.compat.v1.train.Saver can handle them automatically when the graph structure matches.'
• Restore succeeds only if scope, variable names, shape, and partitioning layout stay compatible.
• 'tf.train.list_variables is the quickest debugging tool for checkpoint mismatches.'
• For new code, prefer modern TensorFlow 2 checkpointing unless legacy partitioned-variable workflows are required.