Practice Now
TensorFlow
Linear Regression
Python
Attribute Matrix
Target Matrix

Use attribute and target matrices for TensorFlow Linear Regression Python

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
TensorFlow & KerasMachine LearningPython

Introduction

Linear regression in TensorFlow involves two core data structures: the attribute matrix (X) containing input features and the target matrix (Y) containing the values to predict. Properly shaping and preparing these matrices is essential — TensorFlow operations expect specific dimensions, and mismatched shapes are the most common source of errors in regression models.

Understanding the Matrices

  • Attribute Matrix (X): Shape (n_samples, n_features). Each row is one observation, each column is a feature.
  • Target Matrix (Y): Shape (n_samples, 1) or (n_samples,). The output variable to predict.
python
1import numpy as np
2
3# Example: predicting house prices
4# Features: square_feet, bedrooms, age
5X = np.array([
6    [1500, 3, 10],
7    [2000, 4, 5],
8    [1200, 2, 20],
9    [1800, 3, 8],
10    [2500, 5, 2]
11], dtype=np.float32)  # Shape: (5, 3)
12
13# Target: price in thousands
14Y = np.array([300, 400, 250, 350, 500], dtype=np.float32)  # Shape: (5,)
15Y = Y.reshape(-1, 1)  # Shape: (5, 1) — required for matrix operations

The simplest approach using tf.keras:

python
1import tensorflow as tf
2import numpy as np
3from sklearn.model_selection import train_test_split
4from sklearn.preprocessing import StandardScaler
5
6# Generate sample data
7np.random.seed(42)
8X = np.random.randn(1000, 3).astype(np.float32)
9true_weights = np.array([[2.0], [3.0], [-1.0]])
10Y = X @ true_weights + 0.5 * np.random.randn(1000, 1).astype(np.float32)
11
12# Split data
13X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2)
14
15# Normalize features
16scaler = StandardScaler()
17X_train = scaler.fit_transform(X_train)
18X_test = scaler.transform(X_test)
19
20# Build model
21model = tf.keras.Sequential([
22    tf.keras.layers.Dense(1, input_shape=(3,))  # Linear: Y = X @ W + b
23])
24
25model.compile(optimizer='adam', loss='mse', metrics=['mae'])
26
27# Train
28history = model.fit(X_train, Y_train, epochs=100, batch_size=32,
29                    validation_split=0.2, verbose=0)
30
31# Evaluate
32loss, mae = model.evaluate(X_test, Y_test)
33print(f"Test MSE: {loss:.4f}, MAE: {mae:.4f}")
34
35# Get learned weights
36weights, bias = model.layers[0].get_weights()
37print(f"Weights: {weights.flatten()}")  # Close to [2.0, 3.0, -1.0]
38print(f"Bias: {bias[0]:.4f}")           # Close to 0.0

Method 2: Low-Level TensorFlow

For understanding the math behind linear regression:

python
1import tensorflow as tf
2import numpy as np
3
4# Data
5X_data = np.random.randn(100, 2).astype(np.float32)
6Y_data = (X_data @ np.array([[3.0], [-2.0]]) + 1.0).astype(np.float32)
7
8# Convert to tensors
9X = tf.constant(X_data)
10Y = tf.constant(Y_data)
11
12# Variables (learnable parameters)
13W = tf.Variable(tf.random.normal([2, 1]))
14b = tf.Variable(tf.zeros([1]))
15
16# Training loop
17learning_rate = 0.1
18for epoch in range(200):
19    with tf.GradientTape() as tape:
20        predictions = tf.matmul(X, W) + b  # Y_hat = X @ W + b
21        loss = tf.reduce_mean(tf.square(predictions - Y))  # MSE
22
23    gradients = tape.gradient(loss, [W, b])
24    W.assign_sub(learning_rate * gradients[0])
25    b.assign_sub(learning_rate * gradients[1])
26
27    if epoch % 50 == 0:
28        print(f"Epoch {epoch}: Loss = {loss.numpy():.4f}")
29
30print(f"Learned W: {W.numpy().flatten()}")  # Close to [3.0, -2.0]
31print(f"Learned b: {b.numpy()[0]:.4f}")     # Close to 1.0

Method 3: Normal Equation (Closed-Form)

For small datasets, solve analytically:

python
1import tensorflow as tf
2import numpy as np
3
4X_data = np.random.randn(100, 2).astype(np.float32)
5Y_data = (X_data @ np.array([[3.0], [-2.0]]) + 1.0).astype(np.float32)
6
7# Add bias column (column of ones)
8X_with_bias = tf.concat([tf.ones([100, 1]), tf.constant(X_data)], axis=1)
9
10# Normal equation: W = (X^T X)^(-1) X^T Y
11XtX = tf.matmul(tf.transpose(X_with_bias), X_with_bias)
12XtY = tf.matmul(tf.transpose(X_with_bias), tf.constant(Y_data))
13W = tf.matmul(tf.linalg.inv(XtX), XtY)
14
15print(f"Weights (bias, w1, w2): {W.numpy().flatten()}")
16# Close to [1.0, 3.0, -2.0]

Loading Data from CSV

python
1import pandas as pd
2import numpy as np
3import tensorflow as tf
4
5# Load CSV
6df = pd.read_csv('housing.csv')
7
8# Separate features and target
9feature_cols = ['square_feet', 'bedrooms', 'age']
10X = df[feature_cols].values.astype(np.float32)
11Y = df['price'].values.astype(np.float32).reshape(-1, 1)
12
13# Using tf.data for efficient batching
14dataset = tf.data.Dataset.from_tensor_slices((X, Y))
15dataset = dataset.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)
16
17# Train with dataset
18model.fit(dataset, epochs=50)

Feature Normalization

Always normalize features before training — linear regression converges much faster with normalized inputs:

python
1from sklearn.preprocessing import StandardScaler
2
3scaler = StandardScaler()
4X_train_scaled = scaler.fit_transform(X_train)  # Mean=0, Std=1
5X_test_scaled = scaler.transform(X_test)         # Same transform
6
7# Or with TensorFlow's Normalization layer
8normalizer = tf.keras.layers.Normalization(axis=-1)
9normalizer.adapt(X_train)
10
11model = tf.keras.Sequential([
12    normalizer,
13    tf.keras.layers.Dense(1)
14])

Common Pitfalls

  • Shape mismatch: X must be (n_samples, n_features) and Y must be (n_samples, 1). A common error is Y being (n_samples,) — use Y.reshape(-1, 1).
  • Feature scaling: Without normalization, features with large ranges dominate the loss. Always standardize (zero mean, unit variance) or normalize (0-1 range) before training.
  • dtype mismatch: TensorFlow defaults to float32. Ensure X and Y are both float32, not float64 (NumPy's default).
  • Learning rate too high: For gradient descent, a high learning rate causes divergence (loss goes to infinity). Start with 0.001 and adjust.
  • Overfitting with many features: Linear regression with many features relative to samples can overfit. Use L1/L2 regularization: tf.keras.layers.Dense(1, kernel_regularizer=tf.keras.regularizers.l2(0.01)).

Summary

  • The attribute matrix X has shape (n_samples, n_features), target Y has shape (n_samples, 1)
  • Use tf.keras.layers.Dense(1) for the simplest Keras linear regression
  • Use tf.GradientTape for manual gradient descent (educational purposes)
  • Use the normal equation W = (X^T X)^(-1) X^T Y for small datasets
  • Always normalize features before training for faster convergence
  • Ensure both X and Y are float32 and properly shaped
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
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.