TensorFlow Operations for Linear Algebra and Tensor Manipulation

Introduction: Importance of Linear Algebra in Deep Learning Architectures

Linear algebra is the backbone of deep learning, powering the computations that drive neural network architectures. Operations like matrix multiplication, eigenvalue decomposition, and tensor transformations are critical for tasks such as feature extraction, weight updates, and data processing. In deep learning, layers like fully connected networks rely on matrix multiplications (tf.linalg.matmul) to transform inputs, while operations like singular value decomposition (tf.linalg.svd) are used in techniques like principal component analysis for dimensionality reduction. Tensor manipulations, such as reshaping and slicing, are essential for handling multidimensional data like images or sequences, ensuring compatibility across layers. Understanding these operations in TensorFlow enables developers to build efficient, scalable, and robust deep learning models, optimizing both performance and accuracy.

This guide explores TensorFlow's tf.linalg module for linear algebra, tensor manipulation functions like squeeze and expand_dims, common dimension mismatch issues, a real-world image processing example, and a custom Keras layer leveraging these operations.

1. Linear Algebra Operations (tf.linalg)

TensorFlow's tf.linalg module offers a wide range of linear algebra operations for matrix and tensor computations. Below is a detailed list of key operations, their descriptions, and examples.

Matrix Operations

• tf.linalg.matmul: Performs matrix multiplication between two tensors.
  • Input: Two tensors of shape [..., M, N] and [..., N, P].
  • Output: Tensor of shape [..., M, P].
  • Example:

    import tensorflow as tf
    
    a = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
    b = tf.constant([[5, 6], [7, 8]], dtype=tf.float32)
    result = tf.linalg.matmul(a, b)
    # Result: [[19, 22], [43, 50]]
    

• tf.linalg.inv: Computes the inverse of a square matrix.
  • Input: A square tensor of shape [..., N, N].
  • Output: Tensor of the same shape, representing the inverse.
  • Example:

    matrix = tf.constant([[4, 3], [3, 2]], dtype=tf.float32)
    inverse = tf.linalg.inv(matrix)
    # Result: [[-2, 3], [3, -4]]
    

• tf.linalg.det: Computes the determinant of a square matrix.
  • Input: A square tensor of shape [..., N, N].
  • Output: Tensor of shape [...], containing the determinant.
  • Example:

    matrix = tf.constant([[4, 3], [3, 2]], dtype=tf.float32)
    det = tf.linalg.det(matrix)
    # Result: -1.0
    

• tf.linalg.trace: Computes the trace (sum of diagonal elements) of a matrix.
  • Input: A tensor of shape [..., N, N].
  • Output: Tensor of shape [...].
  • Example:

    matrix = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
    trace = tf.linalg.trace(matrix)
    # Result: 5.0
    

• tf.linalg.transpose: Transposes the last two dimensions of a tensor.
  • Input: A tensor of shape [..., M, N].
  • Output: Tensor of shape [..., N, M].
  • Example:

    matrix = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
    transposed = tf.linalg.transpose(matrix)
    # Result: [[1, 3], [2, 4]]
    

Decomposition Operations

• tf.linalg.cholesky: Computes the Cholesky decomposition of a symmetric positive-definite matrix.
  • Input: A tensor of shape [..., N, N].
  • Output: Tensor of the same shape, representing the lower triangular factor.
  • Example:

    matrix = tf.constant([[4, 2], [2, 2]], dtype=tf.float32)
    cholesky = tf.linalg.cholesky(matrix)
    # Result: Lower triangular matrix
    

• tf.linalg.svd: Computes the Singular Value Decomposition (SVD) of a matrix.
  • Input: A tensor of shape [..., M, N].
  • Output: Singular values, left singular vectors, and right singular vectors.
  • Example:

    matrix = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
    s, u, v = tf.linalg.svd(matrix)
    # s: Singular values, u: Left singular vectors, v: Right singular vectors
    

• tf.linalg.eigh: Computes the eigenvalues and eigenvectors of a symmetric/hermitian matrix.
  • Input: A tensor of shape [..., N, N].
  • Output: Eigenvalues and eigenvectors.
  • Example:

    matrix = tf.constant([[4, 2], [2, 3]], dtype=tf.float32)
    eigenvalues, eigenvectors = tf.linalg.eigh(matrix)
    # Result: Eigenvalues and corresponding eigenvectors
    

Other Linear Algebra Operations

• tf.linalg.norm: Computes the norm of a tensor (e.g., Frobenius norm for matrices).
  • Input: A tensor and optional norm type (fro for Frobenius, euclidean, etc.).
  • Output: Tensor of shape [...].
  • Example:

    matrix = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
    norm = tf.linalg.norm(matrix)
    # Result: Frobenius norm
    

• tf.linalg.solve: Solves a system of linear equations Ax = b.
  • Input: Matrix A of shape [..., N, N] and tensor b of shape [..., N, K].
  • Output: Tensor of shape [..., N, K].
  • Example:

    A = tf.constant([[3, 1], [1, 3]], dtype=tf.float32)
    b = tf.constant([[9], [15]], dtype=tf.float32)
    x = tf.linalg.solve(A, b)
    # Result: Solution to Ax = b
    

• tf.linalg.band_part: Extracts the banded part of a matrix (e.g., diagonal, upper, or lower triangular).
  • Input: A tensor, number of lower and upper diagonals to keep.
  • Output: Tensor with specified bands.
  • Example:

    matrix = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=tf.float32)
    banded = tf.linalg.band_part(matrix, 1, 1)  # Keep main diagonal and one off-diagonal
    # Result: Banded matrix
    

2. Tensor Manipulation Operations

TensorFlow provides a variety of functions to manipulate the shape, structure, and content of tensors. Below are key operations, including squeeze, expand_dims (equivalent to unsqueeze), and others, along with a real-world example of manipulating an RGB image.

Shape Manipulation

• tf.squeeze: Removes dimensions of size 1 from a tensor.
  • Input: A tensor and optional axis to squeeze.
  • Output: Tensor with singleton dimensions removed.
  • Example:

    tensor = tf.constant([[[1], [2]], [[3], [4]]], dtype=tf.float32)  # Shape: (2, 2, 1)
    squeezed = tf.squeeze(tensor)  # Shape: (2, 2)
    # Result: [[1, 2], [3, 4]]
    

• tf.expand_dims: Adds a dimension of size 1 at the specified axis (equivalent to unsqueeze in PyTorch).
  • Input: A tensor and axis to expand.
  • Output: Tensor with an additional dimension.
  • Example:

    tensor = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)  # Shape: (2, 2)
    expanded = tf.expand_dims(tensor, axis=1)  # Shape: (2, 1, 2)
    # Result: [[[1, 2]], [[3, 4]]]
    

• tf.reshape: Reshapes a tensor to a new shape, preserving the total number of elements.
  • Input: A tensor and the target shape.
  • Output: Tensor with the new shape.
  • Example:

    tensor = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)  # Shape: (2, 2)
    reshaped = tf.reshape(tensor, [4])  # Shape: (4,)
    # Result: [1, 2, 3, 4]
    

• tf.transpose: Permutes the dimensions of a tensor according to the specified permutation.
  • Input: A tensor and permutation of axes.
  • Output: Tensor with permuted dimensions.
  • Example:

    tensor = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)  # Shape: (2, 2)
    transposed = tf.transpose(tensor, perm=[1, 0])  # Shape: (2, 2)
    # Result: [[1, 3], [2, 4]]
    

Slicing and Joining

• tf.concat: Concatenates a list of tensors along a specified axis.
  • Input: List of tensors and axis.
  • Output: Concatenated tensor.
  • Example:

    a = tf.constant([[1, 2]], dtype=tf.float32)
    b = tf.constant([[3, 4]], dtype=tf.float32)
    concat = tf.concat([a, b], axis=0)  # Shape: (2, 2)
    # Result: [[1, 2], [3, 4]]
    

• tf.stack: Stacks a list of tensors along a new axis.
  • Input: List of tensors and axis.
  • Output: Stacked tensor with an additional dimension.
  • Example:

    a = tf.constant([1, 2], dtype=tf.float32)
    b = tf.constant([3, 4], dtype=tf.float32)
    stacked = tf.stack([a, b], axis=0)  # Shape: (2, 2)
    # Result: [[1, 2], [3, 4]]
    

• tf.slice: Extracts a slice from a tensor.
  • Input: A tensor, begin indices, and size of the slice.
  • Output: Sliced tensor.
  • Example:

    tensor = tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32)
    sliced = tf.slice(tensor, [0, 1], [2, 2])  # Shape: (2, 2)
    # Result: [[2, 3], [5, 6]]
    

Other Manipulation Operations

• tf.tile: Repeats a tensor a specified number of times.
  • Input: A tensor and multiples for each dimension.
  • Output: Tiled tensor.
  • Example:

    tensor = tf.constant([[1, 2]], dtype=tf.float32)  # Shape: (1, 2)
    tiled = tf.tile(tensor, [2, 1])  # Shape: (2, 2)
    # Result: [[1, 2], [1, 2]]
    

• tf.pad: Pads a tensor with zeros or other values.
  • Input: A tensor and padding configuration.
  • Output: Padded tensor.
  • Example:

    tensor = tf.constant([[1, 2]], dtype=tf.float32)  # Shape: (1, 2)
    padded = tf.pad(tensor, [[1, 1], [0, 0]])  # Shape: (3, 2)
    # Result: [[0, 0], [1, 2], [0, 0]]
    

• tf.reverse: Reverses specific dimensions of a tensor.
  • Input: A tensor and axes to reverse.
  • Output: Tensor with reversed dimensions.
  • Example:

    tensor = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)  # Shape: (2, 2)
    reversed_tensor = tf.reverse(tensor, [0])  # Reverse along axis 0
    # Result: [[3, 4], [1, 2]]
    

Real-World Example: Manipulating an RGB Image

Below is an example demonstrating tensor manipulation operations on an RGB image (3 channels) using TensorFlow. The operations include resizing, applying a filter (convolution), slicing a region, and flipping the image. This simulates common image processing tasks in computer vision applications.

import tensorflow as tf
import numpy as np

# Create a synthetic RGB image (height=64, width=64, channels=3)
image = tf.random.uniform((64, 64, 3), minval=0, maxval=255, dtype=tf.float32)

# 1. Resize the image to 32x32
resized_image = tf.image.resize(image, [32, 32], method='bilinear')
# Shape: (32, 32, 3)

# 2. Apply a 3x3 Gaussian blur filter (convolution)
# Define a 3x3 Gaussian kernel
kernel = tf.constant([
    [1, 2, 1],
    [2, 4, 2],
    [1, 2, 1]
], dtype=tf.float32) / 16.0
kernel = tf.expand_dims(tf.expand_dims(kernel, axis=-1), axis=-1)  # Shape: (3, 3, 1, 1)
kernel = tf.tile(kernel, [1, 1, 3, 1])  # Replicate for 3 channels
filtered_image = tf.nn.conv2d(
    input=tf.expand_dims(resized_image, axis=0),  # Add batch dimension: (1, 32, 32, 3)
    filters=kernel,
    strides=[1, 1, 1, 1],
    padding='SAME'
)
filtered_image = tf.squeeze(filtered_image, axis=0)  # Remove batch dimension: (32, 32, 3)

# 3. Slice a 16x16 region from the center
start_h = (32 - 16) // 2  # Start at height index 8
start_w = (32 - 16) // 2  # Start at width index 8
sliced_image = tf.slice(filtered_image, [start_h, start_w, 0], [16, 16, 3])
# Shape: (16, 16, 3)

# 4. Flip the sliced image horizontally
flipped_image = tf.reverse(sliced_image, axis=[1])
# Shape: (16, 16, 3)

# Print shapes to verify
print("Original image shape:", image.shape)
print("Resized image shape:", resized_image.shape)
print("Filtered image shape:", filtered_image.shape)
print("Sliced image shape:", sliced_image.shape)
print("Flipped image shape:", flipped_image.shape)
# Example output:
# Original image shape: (64, 64, 3)
# Resized image shape: (32, 32, 3)
# Filtered image shape: (32, 32, 3)
# Sliced image shape: (16, 16, 3)
# Flipped image shape: (16, 16, 3)

Explanation: This example processes an RGB image using:

• tf.image.resize: Resizes the image to 32x32 pixels, preserving the 3 channels.
• tf.expand_dims and tf.squeeze: Manage batch dimensions for convolution compatibility.
• tf.nn.conv2d: Applies a Gaussian blur by convolving with a 3x3 kernel, replicated across channels using tf.tile.
• tf.slice: Extracts a 16x16 region from the center of the image.
• tf.reverse: Flips the sliced image horizontally along the width axis.
• Use Case: These operations are typical in computer vision pipelines, such as preprocessing images for convolutional neural networks (CNNs).

3. Common Dimension Mismatch Issues and Examples

Dimension or size mismatches are common errors in TensorFlow when performing linear algebra or tensor manipulation operations. These errors occur when the shapes of input tensors do not meet the operation's requirements. Below are common scenarios, examples of errors, and how to fix them.

Matrix Multiplication (tf.linalg.matmul)

• Issue: For matrix multiplication, the inner dimensions of the two tensors must match (i.e., for shapes [..., M, N] and [..., N, P], the N dimensions must be equal).
• Error Example:

  a = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)  # Shape: (2, 2)
  b = tf.constant([[5, 6, 7], [8, 9, 10]], dtype=tf.float32)  # Shape: (2, 3)
  result = tf.linalg.matmul(a, b)  # Error: Incompatible shapes (2 != 3)
  # InvalidArgumentError: Matrix size-incompatible: In[0]: [2,2], In[1]: [2,3]
  

• Fix: Ensure the inner dimensions match by reshaping or transposing one of the tensors.

  b = tf.constant([[5, 6], [7, 8], [9, 10]], dtype=tf.float32)  # Shape: (3, 2)
  b_transposed = tf.linalg.transpose(b)  # Shape: (2, 3)
  result = tf.linalg.matmul(a, b_transposed)  # Shape: (2, 3)
  # Result: [[19, 25, 31], [43, 57, 71]]
  

Linear System Solving (tf.linalg.solve)

• Issue: The matrix A must be square ([..., N, N]), and the tensor b must have a compatible shape ([..., N, K]).
• Error Example:

  A = tf.constant([[3, 1], [1, 3]], dtype=tf.float32)  # Shape: (2, 2)
  b = tf.constant([9, 15, 21], dtype=tf.float32)  # Shape: (3,)
  x = tf.linalg.solve(A, b)  # Error: Incompatible shapes
  # InvalidArgumentError: Incompatible shapes: [2,2] vs. [3]
  

• Fix: Reshape b to match the expected shape [..., N, K] (e.g., [2, 1]).

  b = tf.constant([[9], [15]], dtype=tf.float32)  # Shape: (2, 1)
  x = tf.linalg.solve(A, b)  # Shape: (2, 1)
  # Result: Solution to Ax = b
  

Concatenation (tf.concat)

• Issue: All tensors must have the same shape except for the concatenation axis.
• Error Example:

  a = tf.constant([[1, 2]], dtype=tf.float32)  # Shape: (1, 2)
  b = tf.constant([[3, 4, 5]], dtype=tf.float32)  # Shape: (1, 3)
  concat = tf.concat([a, b], axis=0)  # Error: Incompatible shapes
  # InvalidArgumentError: ConcatOp : Dimensions of inputs should match
  

• Fix: Ensure the shapes match except along the concatenation axis.

  b = tf.constant([[3, 4]], dtype=tf.float32)  # Shape: (1, 2)
  concat = tf.concat([a, b], axis=0)  # Shape: (2, 2)
  # Result: [[1, 2], [3, 4]]
  

Squeeze and Expand Dimensions (tf.squeeze, tf.expand_dims)

• Issue: tf.squeeze can only remove dimensions of size 1, and attempting to squeeze a non-singleton dimension causes an error.

  tensor = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)  # Shape: (2, 2)
  squeezed = tf.squeeze(tensor, axis=0)  # Error: Cannot squeeze dimension with size 2
  # InvalidArgumentError: Can not squeeze dim[0], expected a dimension of 1, got 2
  

  Fix: Only squeeze dimensions with size 1.

  tensor = tf.constant([[[1], [2]]], dtype=tf.float32)  # Shape: (1, 2, 1)
  squeezed = tf.squeeze(tensor, axis=0)  # Shape: (2, 1)
  # Result: [[1], [2]]
  

• Issue: tf.expand_dims requires a valid axis within the tensor’s rank range.

  tensor = tf.constant([[1, 2]], dtype=tf.float32)  # Shape: (1, 2)
  expanded = tf.expand_dims(tensor, axis=3)  # Error: Axis out of range
  # InvalidArgumentError: axis 3 is out of bounds for tensor with rank 2
  

  Fix: Choose a valid axis (from -rank-1 to rank).

  expanded = tf.expand_dims(tensor, axis=1)  # Shape: (1, 1, 2)
  # Result: [[[1, 2]]]
  

Reshaping (tf.reshape)

• Issue: The new shape must have the same total number of elements as the original tensor.

  tensor = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)  # Shape: (2, 2), 4 elements
  reshaped = tf.reshape(tensor, [2, 3])  # Error: 2*3 = 6, does not match 4
  # InvalidArgumentError: Input to reshape is a tensor with 4 values, but the requested shape has 6
  

  Fix: Ensure the new shape matches the total number of elements.

  reshaped = tf.reshape(tensor, [4, 1])  # Shape: (4, 1)
  # Result: [[1], [2], [3], [4]]
  

Best Practices to Avoid Mismatch Issues

• Check Shapes: Use tensor.shape or tf.shape(tensor) to verify tensor shapes before operations.
• Broadcasting: Be cautious with broadcasting, as TensorFlow may implicitly broadcast tensors, leading to unexpected results. Explicitly reshape or expand dimensions if needed.
• Debugging: Run operations in eager mode to catch shape errors immediately, as graph mode may defer errors to runtime.
• Shape Inference: For dynamic shapes, use tf.ensure_shape to enforce expected shapes during execution.

4. Example: Custom Tensor Layer with Linear Algebra Operations

Below is an example of a custom Keras layer that uses linear algebra operations (tf.linalg.matmul, tf.linalg.inv) and tensor manipulation functions (tf.expand_dims, tf.squeeze) to perform a transformation on input data. This layer computes a transformation matrix based on the input, applies it via matrix multiplication, and handles shape manipulations to ensure compatibility.

Custom Layer: Linear Transformation with Inverse Adjustment

This custom layer, InverseTransformLayer, takes an input tensor, constructs a transformation matrix (e.g., a covariance-like matrix), computes its inverse, and applies the inverse to the input via matrix multiplication. The layer uses tf.expand_dims and tf.squeeze to handle batch dimensions and ensure shape compatibility.

import tensorflow as tf
from tensorflow.keras import layers

class InverseTransformLayer(layers.Layer):
    def __init__(self, output_dim, **kwargs):
        super(InverseTransformLayer, self).__init__(**kwargs)
        self.output_dim = output_dim

    def build(self, input_shape):
        # Initialize a learnable weight matrix for transformation
        self.kernel = self.add_weight(
            name='kernel',
            shape=(input_shape[-1], self.output_dim),
            initializer='glorot_uniform',
            trainable=True
        )

    def call(self, inputs):
        # Ensure inputs have batch dimension
        input_shape = tf.shape(inputs)
        if len(input_shape) == 2:
            inputs = tf.expand_dims(inputs, axis=0)  # Add batch dimension if needed

        # Compute a transformation matrix (e.g., covariance-like matrix)
        # Shape: [batch_size, input_dim, input_dim]
        transform_matrix = tf.matmul(inputs, inputs, transpose_b=True)
        
        # Add a small value to diagonal for numerical stability
        eye = tf.eye(tf.shape(transform_matrix)[-1], batch_shape=[tf.shape(transform_matrix)[0]])
        transform_matrix = transform_matrix + 1e-6 * eye

        # Compute inverse of transformation matrix
        # Shape: [batch_size, input_dim, input_dim]
        inverse_matrix = tf.linalg.inv(transform_matrix)

        # Apply transformation: multiply input by inverse matrix and kernel
        # Shape: [batch_size, input_dim, output_dim]
        intermediate = tf.matmul(inputs, self.kernel)
        
        # Shape: [batch_size, input_dim, output_dim]
        output = tf.matmul(inverse_matrix, intermediate)

        # Squeeze batch dimension if inputs were originally 2D
        if len(input_shape) == 2:
            output = tf.squeeze(output, axis=0)

        return output

    def compute_output_shape(self, input_shape):
        if len(input_shape) == 2:
            return (input_shape[0], self.output_dim)
        return (input_shape[0], input_shape[1], self.output_dim)

# Example usage in a model
input_dim = 4
output_dim = 2
batch_size = 32

# Create a simple model with the custom layer
inputs = tf.keras.Input(shape=(input_dim,))
outputs = InverseTransformLayer(output_dim=output_dim)(inputs)
model = tf.keras.Model(inputs, outputs)

# Test with sample data
sample_input = tf.random.normal((batch_size, input_dim))
output = model(sample_input)
print("Input shape:", sample_input.shape)
print("Output shape:", output.shape)
# Example output: Input shape: (32, 4), Output shape: (32, 2)

Explanation of the Custom Layer

• Purpose: The InverseTransformLayer applies a linear transformation to the input, adjusted by the inverse of a dynamically computed transformation matrix (e.g., a covariance-like matrix derived from the input).
• Linear Algebra Operations:
  • tf.matmul(inputs, inputs, transpose_b=True): Computes a covariance-like matrix.
  • tf.linalg.inv: Computes the inverse of the transformation matrix for adjustment.
  • tf.matmul: Applies the inverse matrix and a learnable kernel to the input.
• Tensor Manipulation:
  • tf.expand_dims: Adds a batch dimension to handle both batched and non-batched inputs.
  • tf.squeeze: Removes the batch dimension if the input was originally 2D.
• Shape Handling: The layer checks the input shape and adjusts the output accordingly, ensuring compatibility with both 2D (single sample) and 3D (batched) inputs.
• Numerical Stability: A small value (1e-6) is added to the diagonal of the transformation matrix to ensure tf.linalg.inv is numerically stable.
• Use Case: This layer could be used in applications like data normalization, feature transformation, or as part of a neural network for tasks requiring adaptive linear transformations.

Potential Dimension Mismatch in the Custom Layer

• Issue: If the input shape is incompatible with the kernel shape, tf.matmul will fail.
• Mitigation: The build method ensures the kernel shape matches the input dimension (input_shape[-1]).
• Issue: The inverse operation (tf.linalg.inv) requires a square matrix.
• Mitigation: The transformation matrix is constructed as a square matrix using tf.matmul(inputs, inputs, transpose_b=True).
• Issue: Mismatches in batch dimensions.
• Mitigation: The layer uses tf.expand_dims and tf.squeeze to handle batch dimensions dynamically.

5. Notes and Best Practices

• Data Types: Ensure consistent data types (e.g., tf.float32, tf.float64) to avoid type mismatch errors.
• Batching: Most tf.linalg operations support batched inputs (tensors with leading dimensions), enabling parallel computation.
• Eager Execution: TensorFlow operations work in both eager and graph modes. For dynamic shapes, eager execution is often more intuitive.
• Performance: Use GPU/TPU acceleration for large-scale linear algebra operations to improve performance.
• Numerical Stability: For operations like tf.linalg.inv or tf.linalg.cholesky, ensure inputs are well-conditioned to avoid numerical errors.
• Custom Layers: When designing custom layers, always validate input shapes and handle batch dimensions explicitly to ensure compatibility with various input formats.

6. Example: Combining Linear Algebra and Tensor Manipulation

Below is a complete example that combines matrix multiplication, tensor reshaping, and squeezing, with checks to avoid dimension mismatches:

import tensorflow as tf

# Create two matrices
A = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)  # Shape: (2, 2)
B = tf.constant([[5, 6], [7, 8]], dtype=tf.float32)  # Shape: (2, 2)

# Check shapes before matrix multiplication
assert A.shape[-1] == B.shape[-2], "Inner dimensions must match for matmul"
C = tf.linalg.matmul(A, B)  # Shape: (2, 2)

# Reshape to a 1D tensor
reshaped = tf.reshape(C, [-1])  # Shape: (4,)

# Add a singleton dimension
expanded = tf.expand_dims(reshaped, axis=0)  # Shape: (1, 4)

# Squeeze the singleton dimension
squeezed = tf.squeeze(expanded, axis=0)  # Shape: (4,)

print("Matrix multiplication result:\n", C)
print("Reshaped tensor:", reshaped)
print("Expanded tensor:", expanded)
print("Squeezed tensor:", squeezed)

This guide covers the core linear algebra and tensor manipulation operations in TensorFlow, their importance in deep learning, common dimension mismatch issues, a practical image processing example, and a custom layer leveraging these operations. For more details, refer to the TensorFlow documentation.