Welcome back, intrepid explorer! I'm thrilled to see you for the second lesson in our JAX in Action: Neural Networks from Scratch course. In our previous transmission, we laid the critical groundwork: we prepared our XOR dataset and skillfully crafted the initialize_mlp_params function to set up our network's weights and biases using the Xavier/Glorot method. Our parameters are now neatly organized in a PyTree, ready for action!

Today, we're shifting gears to something truly dynamic: implementing forward propagation. This is the very heart of how a neural network makes predictions. We'll build a function that takes our carefully prepared inputs and parameters and channels the data through the network's layers and activation functions to produce an output. By the end of this lesson, you'll have a working JAX function that performs the complete forward pass for our Multi-Layer Perceptron (MLP), and you'll see it generate its very first (untrained) predictions for the XOR problem!

Now, let's translate this understanding into a JAX function. We'll create mlp_forward_pass, which takes the list of parameters (our PyTree of weights and biases for each layer) and an input x_input. It will then guide the input data through the network.

@jax.jit
def mlp_forward_pass(params_list, x_input):
    """
    Performs the forward pass for a Multi-Layer Perceptron.
    
    Args:
        params_list: List of layer parameters (dicts with 'w' and 'b').
                     Each dict contains 'w' (weights) and 'b' (biases).
        x_input: Input data (can be a single sample or a batch of samples).
    
    Returns:
        Output predictions from the MLP.
    """
    # Start with the input data as the initial "activations"
    activations = x_input
    
    # Process all layers except the last one (hidden layers)
    for i in range(len(params_list) - 1):
        layer_params = params_list[i]
        weights = layer_params['w']
        biases = layer_params['b']
        
        # Affine transformation: z = activations @ weights + biases
        z = jnp.dot(activations, weights) + biases
        
        # Apply sigmoid activation function
        activations = jax.nn.sigmoid(z)
        
    # Process the output layer
    # The last item in params_list corresponds to the output layer
    output_layer_params = params_list[-1]
    output_weights = output_layer_params['w']
    output_biases = output_layer_params['b']
    
    # Affine transformation for the output layer
    z_out = jnp.dot(activations, output_weights) + output_biases
    
    # Apply sigmoid activation for the output layer (common for binary classification)
    output_predictions = jax.nn.sigmoid(z_out)
    
    return output_predictions

Let's break this down:

• We initialize activations with the x_input. This variable will hold the output of the current layer, which becomes the input to the next.
• We loop through params_list up to the second-to-last element. These are our hidden layers.
  • Inside the loop, we extract the weights and biases for the current layer.
  • We perform the affine transformation: z = jnp.dot(activations, weights) + biases.
  • We apply the jax.nn.sigmoid activation function to z and update activations.
• After the loop, we handle the output layer separately using params_list[-1]. This is good practice, as output layers sometimes have different activation functions (though here, we use sigmoid again, suitable for binary classification like XOR).
• Finally, the function returns output_predictions.
• The @jax.jit decorator just-in-time compiles this function for optimized performance, making it run super fast, especially when we call it repeatedly during training!

Before we can see our mlp_forward_pass function in action, we need two things: our input data (the XOR examples) and the initialized parameters for our network. We prepared the XOR data in the previous lesson, and we also developed a function, initialize_mlp_params, to create the network's weights and biases.

Let's bring in the initialize_mlp_params function. You'll recall its detailed construction from our last lesson; we include it here as it's essential for our current task of running the forward pass.

import jax
import jax.numpy as jnp

# (This function was developed in the previous lesson)
def initialize_mlp_params(layer_sizes, parent_key):
    """Initialize MLP parameters"""
    params_list = []
    current_key = parent_key
    
    for i in range(len(layer_sizes) - 1):
        input_dim, output_dim = layer_sizes[i], layer_sizes[i+1]
        # Split key for deterministic random number generation
        current_key, w_key, b_key = jax.random.split(current_key, 3)
        
        # Xavier/Glorot initialization for weights
        limit = jnp.sqrt(6 / (input_dim + output_dim))
        weights = jax.random.uniform(w_key, (input_dim, output_dim), 
                                     minval=-limit, maxval=limit)
        # Biases initialized to zero
        biases = jnp.zeros((output_dim,))
        
        params_list.append({'w': weights, 'b': biases})
        
    return params_list

# Now, let's set up our specific XOR problem data and MLP structure
# XOR input features
xor_X = jnp.array([[0., 0.], [0., 1.], [1., 0.], [1., 1.]], dtype=jnp.float32)

# Define MLP layer sizes: [input_size, hidden_size, output_size]
# For XOR: 2 input features, 3 neurons in hidden layer, 1 output neuron
layer_sizes = [2, 3, 1]

# Generate a JAX PRNG key for parameter initialization
prng_key = jax.random.key(0)

# Initialize the MLP parameters
mlp_params = initialize_mlp_params(layer_sizes, prng_key)

Here, we've defined our xor_X data. Then, we specified layer_sizes for an MLP with an input layer of 2 neurons, a hidden layer of 3 neurons, and an output layer of 1 neuron. Finally, we used jax.random.key(0) to create a PRNG key and called initialize_mlp_params to get our mlp_params PyTree. With our inputs and parameters ready, we can now perform the forward pass!

It's time for the exciting part: let's use our mlp_forward_pass function with the xor_X data and the mlp_params we just prepared. This will give us the initial predictions of our untrained network.

# Perform the forward pass for all XOR inputs
initial_predictions = mlp_forward_pass(mlp_params, xor_X)

print("\nInitial Predictions:")
print(initial_predictions)
print(f"Shape of predictions: {initial_predictions.shape}")

# Let's also test with a single input sample
# We take the first sample: [0., 0.]
# Note: We use xor_X[0:1, :] to keep it as a 2D array (batch size of 1)
single_sample = xor_X[0:1, :] 
print(f"\nSingle input sample: {single_sample}")

prediction_single = mlp_forward_pass(mlp_params, single_sample)
print(f"Prediction for single sample: {prediction_single}")

When you run this code, you'll see the following output:

Initial Predictions:
[[0.24461032]
 [0.30570248]
 [0.3021773 ]
 [0.3624995 ]]
Shape of predictions: (4, 1)

Single input sample: [[0. 0.]]
Prediction for single sample: [[0.24461032]]

Let's analyze this output:

• The "Initial Predictions" are the network's outputs for each of the four XOR input pairs. Since the network's weights were initialized randomly and it hasn't been trained, these predictions are essentially random values (between 0 and 1 due to the sigmoid activation in the output layer). They don't match the true XOR outputs yet (which should be [[0.], [1.], [1.], [0.]]).
• The Shape of predictions: (4, 1) confirms that our network is producing one output for each of the four input samples, which is exactly what we expect.
• The test with a single_sample ([[0. 0.]]) shows that our mlp_forward_pass function correctly handles inputs with a batch size of 1, producing a single prediction [[0.24461032]]. This demonstrates the flexibility of our implementation.

Seeing these outputs, even if random, is a significant milestone. It means our data is flowing correctly through the network architecture we've defined!

Fantastic work, cosmic coder! You've successfully implemented the forward propagation mechanism for our MLP, allowing it to take inputs and generate predictions. This mlp_forward_pass function is a cornerstone of any neural network, representing how it processes information.

While our network currently makes random guesses, we've built the essential pathway for information flow. In our next lesson, we'll tackle the other half of the learning puzzle: calculating how wrong our predictions are (the loss) and figuring out how to adjust our network's parameters to improve them using backpropagation and JAX's powerful automatic differentiation.