ai_ml / ml frameworks / 02_what_is_tensorflow.md

What is TensorFlow?

5 interview angles 5 min read source

What is TensorFlow?

TensorFlow is Google’s open-source deep learning framework (released 2015). The first widely-adopted production-grade DL framework. Now generally paired with Keras as its high-level API (Keras is the official frontend since TF 2.0).

PyTorch overtook TF in research, but TF remains strong in production, mobile/edge deploy, and Google ecosystem integration (TPUs, GCP, TFX).

The TF 1.x vs TF 2.x story

You need this context because interviews sometimes test for it.

  • TF 1.x (2015-2019) — static graphs. You built a graph (tf.placeholder, tf.Variable), launched a tf.Session, ran session.run(...). Painful to debug. Powerful for production / optimization.
  • TF 2.x (2019+) — eager execution by default (like PyTorch). Keras promoted to first-class. tf.function decorator opt-in to graph mode for performance.

If you see code with tf.Session, tf.placeholder, feed_dict — that’s TF 1.x. Old / legacy.

Core idea

import tensorflow as tf

x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
y = tf.random.normal((2, 2))
z = x @ y + 1                              # eager — runs immediately
print(z.shape, z.dtype, z.device)

Like PyTorch, TF gives you:

  1. Tensorstf.Tensor (immutable) and tf.Variable (mutable, trainable).
  2. Autogradtf.GradientTape.
  3. Keras APItf.keras.Model, tf.keras.layers.*.
  4. Optimizers / losses — same families.
  5. tf.data.Dataset — pipeline-style data loading.

Keras — the high-level API

The recommended way to write TF models. Three styles:

Sequential (simplest):

from tensorflow import keras
from tensorflow.keras import layers

model = keras.Sequential([
    layers.Dense(128, activation="relu", input_shape=(784,)),
    layers.Dense(10),
])
model.compile(optimizer="adam", loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=["accuracy"])
model.fit(x_train, y_train, epochs=10, validation_data=(x_val, y_val))

Functional API (graphs with multiple inputs/outputs):

inputs = keras.Input(shape=(784,))
x = layers.Dense(128, activation="relu")(inputs)
outputs = layers.Dense(10)(x)
model = keras.Model(inputs, outputs)

Subclassing (PyTorch-like, max flexibility):

class Net(keras.Model):
    def __init__(self):
        super().__init__()
        self.fc1 = layers.Dense(128, activation="relu")
        self.fc2 = layers.Dense(10)

    def call(self, x):                     # 'call', not 'forward'
        return self.fc2(self.fc1(x))

Most production TF code uses Sequential or Functional. Subclassing is rarer.

Autograd — GradientTape

x = tf.Variable(2.0)
with tf.GradientTape() as tape:
    y = x ** 3 + 2 * x
dy_dx = tape.gradient(y, x)
print(dy_dx)   # 3*x^2 + 2 = 14

Compared to PyTorch’s requires_grad + .backward(), GradientTape is more explicit — you wrap the forward pass and ask for gradients on specific variables.

Custom training loop

If model.fit() isn’t flexible enough:

optimizer = tf.keras.optimizers.AdamW(learning_rate=1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

@tf.function   # compiles to a graph for speed
def train_step(xb, yb):
    with tf.GradientTape() as tape:
        logits = model(xb, training=True)
        loss = loss_fn(yb, logits)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

for epoch in range(10):
    for xb, yb in train_ds:
        loss = train_step(xb, yb)

tf.function — graph mode

@tf.function traces a Python function into a TF graph (XLA-compatible). First call is slow (tracing); subsequent calls run a compiled graph. Big speedup for tight loops on GPU/TPU. Gotcha: side effects (print, mutations) inside @tf.function behave unintuitively because the function may only run once during tracing.

@tf.function
def f(x):
    print("tracing")          # only prints once, during tracing
    tf.print("running", x)    # prints every call
    return x * 2

tf.data.Dataset — input pipelines

ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
ds = (
    ds.shuffle(10_000)
      .batch(64)
      .map(augment, num_parallel_calls=tf.data.AUTOTUNE)
      .prefetch(tf.data.AUTOTUNE)
)

Composable, lazy, parallelized. Designed to overlap data prep with GPU/TPU compute.

Key concepts

Concept What it is
tf.Tensor immutable n-d array
tf.Variable mutable, trainable tensor
tf.keras.Model base class for Keras models (.fit, .compile, .save, .summary)
tf.keras.layers.Layer base class for layers (analog to nn.Module)
tf.GradientTape autograd context — records ops for gradient computation
@tf.function trace Python → TF graph for performance
tf.data.Dataset input pipeline
tf.distribute.Strategy distributed training (mirrored, multi-worker, TPU)
Eager vs graph TF 2 eager by default; graph mode via @tf.function

The TensorFlow ecosystem

  • Keras — high-level API (now also has Keras 3, multi-backend).
  • TensorFlow Lite (TFLite) — mobile / edge / microcontrollers.
  • TensorFlow.js — browser / Node.
  • TensorFlow Serving — production model serving.
  • TFX (TensorFlow Extended) — production ML pipelines (data validation, transform, training, serving).
  • TensorBoard — training visualization (best-in-class; works with PyTorch too).
  • TPU support — TF is the canonical TPU framework (PyTorch/XLA exists but TF is more mature here).
  • Keras 3 — multi-backend (runs on TF, JAX, or PyTorch).

Production strengths

  • TFLite for mobile / edge — better story than PyTorch Mobile historically.
  • TensorFlow.js for browser inference.
  • TPU — Google’s accelerators, first-class TF support.
  • TF Serving — battle-tested model server, mature gRPC API.
  • TFX — end-to-end pipelines used at Google scale.

Common gotchas

  • TF 1.x code in tutorialstf.Session, tf.placeholder, tf.global_variables_initializer is legacy. Use TF 2.x style.
  • @tf.function side effectsprint() runs once during tracing, not every call. Use tf.print() for runtime printing.
  • Variable creation inside @tf.function — creating tf.Variable inside a traced function raises; create them outside.
  • Eager-only ops — some Python operations don’t trace into graphs; you’ll get retracing warnings (slow).
  • training=True/False — must be passed explicitly for layers with different train/eval behavior (Dropout, BatchNorm). Keras .fit() handles this; custom loops don’t.
  • CPU/GPU placement — TF auto-places ops; force with with tf.device("/GPU:0"): if needed.

Keras 3 — the multi-backend pivot

Keras 3 (released 2023) runs on TensorFlow, JAX, or PyTorch as a backend. Same keras.Model / keras.layers API, choose backend via env var:

import os
os.environ["KERAS_BACKEND"] = "jax"      # or "torch" or "tensorflow"
import keras

This is significant: Keras is becoming framework-agnostic. The “Keras model” you write today can run on any backend.

Where to go next

Interview angle

  • “What is TensorFlow?” — Google’s open-source DL framework. TF 2.x is eager-by-default with Keras as the official high-level API. Strong in production, mobile/edge (TFLite), and Google ecosystem (TPU, GCP).
  • “What’s the difference between TF 1.x and TF 2.x?” — 1.x was static graphs (tf.Session, tf.placeholder, feed_dict), painful to debug. 2.x is eager by default with @tf.function for opt-in graph compilation. Keras became the official high-level API.
  • “What does @tf.function do?” — traces a Python function into a TF graph that can be optimized and run faster on GPU/TPU. First call traces; subsequent calls run compiled. Beware: Python side effects only run during tracing.
  • “What’s tf.GradientTape?” — context manager that records operations on watched tensors so you can compute gradients via tape.gradient(loss, variables). TF’s equivalent of PyTorch’s autograd, but explicit instead of automatic.
  • “When would you pick TF over PyTorch?” — when you need TPU training, TFLite for mobile, TF Serving / TFX pipelines, or you’re already in the Google Cloud ecosystem.