!pip install --upgrade -q "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

from functools import partial
import os
import time


os.environ["KERAS_BACKEND"] = "jax"

import jax
from jax.experimental import pallas as pl
import jax.numpy as jnp
import keras

def add_vectors_kernel(x_ref, y_ref, o_ref):
    """Pallas kernel for adding two vectors together."""
    x, y = x_ref[...], y_ref[...]
    o_ref[...] = x + y

@jax.jit
def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array:
    return pl.pallas_call(
        add_vectors_kernel, out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype)
    )(x, y)


add_vectors(jnp.arange(8), jnp.arange(8))

class PallasAddLayer(keras.Layer):
    def call(self, x, y):
        # Reuse the JIT-compiled Pallas function
        return add_vectors(x, y)


layer = PallasAddLayer()

x_data = jnp.arange(8, dtype=jnp.int32)
y_data = jnp.arange(8, dtype=jnp.int32)

layer(x_data, y_data)

class StandardDenseReLU(keras.layers.Layer):
    """Standard Matmul and ReLU implementation using keras.ops."""

    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="glorot_uniform",
            trainable=True,
        )

    def call(self, inputs):
        # The standard implementation performs two separate operations.
        # Each one involves expensive data transfer with the main device memory (HBM).
        # 1. Matmul: inputs (HBM) -> compute -> intermediate (HBM)
        y = keras.ops.matmul(inputs, self.w)
        # 2. ReLU: intermediate (HBM) -> compute -> output (HBM)
        return keras.ops.relu(y)

import jax.numpy as jnp
from jax.experimental import pallas as pl


def matmul_relu_kernel(a_ref, b_ref, c_ref):
    """Pallas kernel for fused matmul + ReLU."""
    # Perform the matrix multiplication on the local tile
    # pl.dot leverages the hardware's Matrix Unit (MXU)
    acc = pl.dot(a_ref[...], b_ref[...])

    # Fusion happens here: apply activation while data is in VMEM
    result = jnp.maximum(acc, 0)

    # Write the final result to the output reference
    c_ref[...] = result

@jax.jit
def fused_matmul(a, b):
    m, k = a.shape
    _, n = b.shape

    # Define tile sizes
    tile_m, tile_n = 128, 128
    assert (
        m % tile_m == 0 and n % tile_n == 0
    ), "Inputs must be multiples of 128 for this demo"

    return pl.pallas_call(
        matmul_relu_kernel,
        # Map output indices to input blocks
        out_shape=jax.ShapeDtypeStruct((m, n), a.dtype),
        in_specs=[
            # For each output tile, we take a slice of A of shape (tile_m, k)
            pl.BlockSpec(
                index_map=lambda i, j: (i, 0), block_shape=(tile_m, k)
            ),  # Matrix A
            # For each output tile, we take a slice of B of shape (k, tile_n)
            pl.BlockSpec(
                index_map=lambda i, j: (0, j), block_shape=(k, tile_n)
            ),  # Matrix B
        ],
        out_specs=pl.BlockSpec(
            index_map=lambda i, j: (i, j), block_shape=(tile_m, tile_n)
        ),  # Matrix C
        grid=(m // tile_m, n // tile_n),
    )(a, b)


fused_matmul(jnp.ones((256, 256)), jnp.ones((256, 256)))

class FusedDense(keras.layers.Layer):
    """Custom Keras layer that applies the fused Dense and ReLU op."""

    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units), initializer="glorot_uniform"
        )

    def call(self, inputs):
        # Dispatch to our Pallas kernel
        return fused_matmul(inputs, self.w.value)


FusedDense(256)(jnp.ones((256, 256)))

# 1. Setup Data
N = 8192  # Large enough to be memory bound
input_data = jnp.ones((N, N), dtype="float32")

# Initialize layers
standard_layer = StandardDenseReLU(units=N)
pallas_layer = FusedDense(units=N)

# Build layers by calling them once
standard_layer(input_data)
pallas_layer(input_data)


def benchmark(layer, x, name, iterations=100):
    # Warm up to ensure JIT compilation is finished
    for _ in range(10):
        layer(x).block_until_ready()

    start_time = time.perf_counter()
    for _ in range(iterations):
        layer(x).block_until_ready()
    end_time = time.perf_counter()

    avg_time = (end_time - start_time) / iterations * 1000  # convert to ms
    print(f"{name} Average Latency: {avg_time:.3f} ms")


# 2. Run Comparison
print(f"Benchmarking Matrix Size: {N}x{N}\n" + "-" * 30)
benchmark(standard_layer, input_data, "Standard Keras (Matmul + ReLU)")
benchmark(pallas_layer, input_data, "Pallas Fused (Matmul + ReLU)")

model = keras.Sequential([FusedDense(256)])
model.compile(optimizer="adam", loss="mse")
model.fit(jnp.ones((256, 256)), jnp.ones((256, 256)))
>>> Linearization failed to produce known values for all output primals. This is
typically caused by attempting to differentiate a function uses an operation
that does not support reverse-mode autodiff.

# 1. Define the wrapper with `custom_vjp` using our original `fused_matmul`.
@jax.custom_vjp
def fused_matmul_trainable(x, w):
    return fused_matmul(x, w)


# 2. Define the Forward Pass
# It must return the output AND "residuals" (data needed for the backward pass)
def fused_matmul_fwd(x, w):
    y = fused_matmul_trainable(x, w)
    # We save inputs x, w and output y for the backward calculation
    return y, (x, w, y)


# 3. Define the Backward Pass
# JAX gives us the residuals and the incoming gradient (g)
def fused_matmul_bwd(residuals, g):
    x, w, y = residuals

    # Calculate the gradient of ReLU: 1 if y > 0, else 0
    # g is the gradient flowing back from the next layer
    grad_relu = g * (y > 0)

    # Standard backprop math for matmul:
    # grad_x = grad_relu @ w.T
    grad_x = jnp.dot(grad_relu, w.T)

    # grad_w = x.T @ grad_relu
    grad_w = jnp.dot(x.T, grad_relu)

    return grad_x, grad_w


# 4. Register the forward and backward functions
fused_matmul_trainable.defvjp(fused_matmul_fwd, fused_matmul_bwd)


class FusedDenseTrainable(FusedDense):
    """Updated layer that contains Pallas forward and backward pass."""

    def call(self, inputs):
        # Dispatch to our trainable Pallas kernel
        return fused_matmul_trainable(inputs, self.w.value)


# Demonstrate trainability on dummy data
model = keras.Sequential([FusedDenseTrainable(256)])
model.compile(optimizer="adam", loss="mse")
model.fit(jnp.ones((256, 256)), jnp.ones((256, 256)), batch_size=128)