Introduction

TensorFlow unit tests should verify tensor behavior quickly and deterministically, without running expensive training loops. Flaky tests usually come from uncontrolled randomness, unstable device assumptions, or loose numeric assertions. A strong test setup focuses on repeatability, small fixtures, and clear failure messages.

Build Deterministic Tensor Tests

Set seeds and keep test inputs small.

1import tensorflow as tf
2
3
4def test_basic_addition():
5    tf.random.set_seed(42)
6    x = tf.constant([[1.0, 2.0]], dtype=tf.float32)
7    y = tf.constant([[3.0, 4.0]], dtype=tf.float32)
8    out = x + y
9    expected = tf.constant([[4.0, 6.0]], dtype=tf.float32)
10    tf.debugging.assert_near(out, expected)

What this validates:

• dtype and shape compatibility,
• numerical output correctness,
• deterministic behavior for fixed inputs.

Prefer unit tests that run in milliseconds and isolate one behavior each.

Test tf.function and Eager Consistency

Many bugs appear only when code is traced by tf.function.

1import tensorflow as tf
2
3@tf.function
4def scale(x):
5    return x * 2.0
6
7
8def test_tf_function_matches_eager():
9    x = tf.constant([1.5, 2.5], dtype=tf.float32)
10    eager_out = x * 2.0
11    graph_out = scale(x)
12    tf.debugging.assert_near(eager_out, graph_out)

This catches differences between eager and traced execution paths early.

Choose Assertion Style for Numeric Stability

Use tolerance-aware assertions for floating-point outputs.

1import numpy as np
2
3
4def test_softmax_close():
5    x = tf.constant([1.0, 2.0, 3.0])
6    y = tf.nn.softmax(x).numpy()
7    expected = np.array([0.09003057, 0.24472847, 0.66524096])
8    np.testing.assert_allclose(y, expected, rtol=1e-6, atol=1e-7)

Use strict equality only for integer or exact symbolic operations.

Run with Pytest and Keep Scope Clear

Run unit tests separately from heavier integration tests.

pytest -q tests/unit
pytest -q tests/unit/test_tensor_ops.py -k "softmax"

A clean test layout helps:

• tests/unit for pure tensor logic,
• tests/integration for data pipelines and training steps,
• tests/e2e for full model workflows.

This keeps feedback loop fast while preserving deeper coverage in slower stages.

CPU and GPU Test Strategy

Do not require GPU for baseline unit tests unless feature is GPU-specific. Keep most unit tests CPU-compatible so CI remains portable.

If GPU-specific kernels must be tested, mark those tests explicitly and run them in dedicated pipelines.

1import tensorflow as tf
2import pytest
3
4
5def has_gpu():
6    return len(tf.config.list_physical_devices("GPU")) > 0
7
8
9@pytest.mark.skipif(not has_gpu(), reason="GPU not available")
10def test_gpu_only_path():
11    x = tf.constant([1.0, 2.0, 3.0])
12    with tf.device("/GPU:0"):
13        y = x * 2.0
14    tf.debugging.assert_near(y, tf.constant([2.0, 4.0, 6.0]))

CI Reliability Checklist

For CI stability:

• log TensorFlow and Python versions,
• pin key dependencies,
• isolate random seeds where stochastic ops are tested,
• fail fast on first critical module failure.

Useful version log command:

python -c "import tensorflow as tf, sys; print(tf.__version__, sys.version)"

When tests suddenly fail after dependency upgrades, version logs make root-cause analysis much faster.

Organize Test Data and Fixtures

Keep synthetic fixtures tiny and local to test modules unless shared across many cases. For shared fixtures, use explicit factory helpers so shape and dtype assumptions stay visible.

Example helper:

1import tensorflow as tf
2
3def make_batch(batch=2, features=3):
4    return tf.ones((batch, features), dtype=tf.float32)

Using helper factories reduces repeated setup code and keeps fixture changes centralized when model interfaces evolve. It also makes new tests easier to review because setup assumptions are standardized. Consistent fixtures reduce flaky CI behavior over time.

Common Pitfalls

• Mixing long training scenarios into unit test stage.
• Using no seed control in stochastic operations.
• Overly broad numeric tolerances that hide regressions.
• Assuming GPU availability in generic CI runners.
• Testing multiple behaviors in one test and getting ambiguous failures.

Summary

• Keep TensorFlow unit tests small, deterministic, and behavior-specific.
• Validate eager and tf.function paths when relevant.
• Use tolerance-aware assertions for floating-point results.
• Separate fast unit tests from heavy integration or training tests.
• Make CI reproducible with version logging and controlled environments.