Tensorflow
Slice Assignment
Machine Learning
Tensor Manipulation
Programming Tutorial

How to do slice assignment in Tensorflow

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Slice assignment in TensorFlow is different from NumPy because regular tensors are immutable. You can still update parts of data efficiently, but you need the right API depending on whether you are working with tf.Variable or pure tensor expressions. This guide covers practical patterns for row updates, block updates, and sparse index updates.

Core Topic Sections

Start with TensorFlow mutability rules

The most important rule is simple:

  1. tf.Tensor is immutable.
  2. tf.Variable is mutable and supports assignment.

If you try NumPy-like direct assignment on tf.Tensor, it will fail. Use assignment methods on tf.Variable, or build a new tensor with scatter operations.

Slice assignment with tf.Variable.assign

For contiguous slices, tf.Variable gives the most direct style.

python
1import tensorflow as tf
2
3x = tf.Variable([[1, 2, 3],
4                 [4, 5, 6],
5                 [7, 8, 9]], dtype=tf.int32)
6
7# Replace row 1
8x[1, :].assign([40, 50, 60])
9
10# Replace a block (rows 0..1, cols 1..2)
11x[0:2, 1:3].assign([[20, 30],
12                    [70, 80]])
13
14print(x.numpy())

This is usually the clearest approach when the tensor should stay mutable during a training or preprocessing step.

Immutable update with tf.tensor_scatter_nd_update

When you need functional style updates without mutable variables, use scatter update to create a new tensor.

python
1import tensorflow as tf
2
3x = tf.constant([[1, 2, 3],
4                 [4, 5, 6],
5                 [7, 8, 9]], dtype=tf.int32)
6
7indices = tf.constant([[0, 1], [2, 2]])
8updates = tf.constant([200, 900])
9
10y = tf.tensor_scatter_nd_update(x, indices, updates)
11
12print("original")
13print(x.numpy())
14print("updated")
15print(y.numpy())

This is ideal in graph-friendly pipelines where immutable transformations are easier to reason about.

Use scatter add or max for accumulation logic

For assignment-like workflows where updates are additive or reduction-based, use dedicated variants:

  1. tf.tensor_scatter_nd_add
  2. tf.tensor_scatter_nd_sub
  3. tf.tensor_scatter_nd_max
  4. tf.tensor_scatter_nd_min

Example with additive updates:

python
1import tensorflow as tf
2
3x = tf.constant([10, 20, 30, 40], dtype=tf.int32)
4indices = tf.constant([[1], [3]])
5updates = tf.constant([5, -10])
6
7y = tf.tensor_scatter_nd_add(x, indices, updates)
8print(y.numpy())  # [10, 25, 30, 30]

This avoids manual slicing and concatenation in many sparse update cases.

Build block updates with slice and concat when needed

For larger rectangular updates on immutable tensors, you can combine slicing and concatenation.

python
1import tensorflow as tf
2
3x = tf.constant([[1., 2., 3., 4.],
4                 [5., 6., 7., 8.],
5                 [9., 10., 11., 12.]])
6
7new_middle = tf.constant([[100., 101.],
8                          [200., 201.],
9                          [300., 301.]])
10
11left = x[:, :1]
12right = x[:, 3:]
13y = tf.concat([left, new_middle, right], axis=1)
14
15print(y.numpy())

This method is verbose but predictable for structured region replacement.

Performance and shape checks

Slice assignment problems are often shape mismatches. Before assignment, confirm:

  1. Target slice shape.
  2. Update tensor shape.
  3. Dtype compatibility.
  4. Device placement expectations.

Quick debug pattern:

python
print("target shape", x[0:2, 1:3].shape)
print("update shape", tf.constant([[1, 2], [3, 4]]).shape)

Small checks prevent hard-to-read runtime errors.

Behavior inside tf.function

Both variable assignment and scatter operations work in tf.function, but mutable state can make execution order harder to follow. Keep update logic explicit and avoid hidden side effects across unrelated functions.

If deterministic behavior matters, return updated tensors from each step rather than mutating deep shared variables.

Choosing the right pattern

Use this practical rule:

  1. Stateful model buffer, prefer tf.Variable.assign.
  2. Pure data transform, prefer scatter update.
  3. Structured block replace, slice plus concat is acceptable.

Consistent style across a codebase improves readability and reduces debugging time.

Common Pitfalls

  • Treating tf.Tensor like NumPy arrays and trying direct in-place assignment.
  • Updating slices with incompatible shapes and getting opaque runtime errors.
  • Mixing mutable and immutable update styles in one pipeline without clear intent.
  • Recomputing large tensors repeatedly when sparse scatter would be cheaper.
  • Using assignment-heavy logic without verifying behavior under tf.function.

Summary

  • TensorFlow slice assignment depends on mutability model.
  • Use tf.Variable.assign for direct mutable updates.
  • Use scatter operations for immutable functional updates.
  • Validate slice and update shapes before writing.
  • Pick one update style per workflow for predictable, maintainable code.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.