Keras.io Documentation Generator

Data augmentation helps prevent overfitting by making the model robust to changes in lighting, cropping, and orientation. Common layers available in keras.layers include:

• RandomFlip(): Flips the image (use with caution on datasets where orientation matters, like MNIST).
• RandomCrop(height, width): Selects a random subset of the image to encourage spatial invariance.
• RandomRotation(factor): Rotates the image by a random angle within a specified range (e.g., (-45 / 360, 45 / 360)).
• Resizing(*size): Used deterministically in evaluation pipelines to ensure consistent input dimensions without adding noise.

# Example augmentation pipeline
random_flip = keras.layers.RandomFlip()
crop = keras.layers.RandomCrop(int(IMAGE_SIZE[0] * 0.9), int(IMAGE_SIZE[1] * 0.9))
rotate = keras.layers.RandomRotation((-45 / 360, 45 / 360))

# Applying to a batch
image_batch = random_flip(image_batch)
image_batch = crop(image_batch)
image_batch = rotate(image_batch)

You can embed a JIT-compiled Pallas function into a Keras layer by calling the function within the layer's call method.

1. Define the Pallas kernel function using references (_ref).
2. Wrap the kernel in pl.pallas_call and decorate it with @jax.jit.
3. Create a keras.Layer subclass and call the JIT-compiled function in call().

from functools import partial
import os
import jax
from jax.experimental import pallas as pl
import jax.numpy as jnp
import keras

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

# 1. Define the kernel
def add_vectors_kernel(x_ref, y_ref, o_ref):
    x, y = x_ref[...], y_ref[...]
    o_ref[...] = x + y

# 2. JIT-compile the Pallas function
@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)

# 3. Embed in a Keras Layer
class PallasAddLayer(keras.Layer):
    def call(self, x, y):
        return add_vectors(x, y)

layer = PallasAddLayer()
layer(jnp.arange(8), jnp.arange(8))

Generative Adversarial Networks (GANs) consist of two models: a generator that maps latent space points to image space, and a discriminator that classifies images as real or fake.

To implement the training loop:

1. Train the discriminator:
   • Sample random points in the latent space and generate fake images via the generator.
   • Combine fake images with real images from the dataset.
   • Train the discriminator to correctly classify generated vs. real images.
2. Train the generator:
   • Sample random points in the latent space and generate fake images.
   • Train the generator to 'fool' the discriminator by attempting to classify the fake images as real.

Note: This implementation requires the TensorFlow backend and is best run on a GPU.

# Example Discriminator Architecture
discriminator = keras.Sequential(
    [
        keras.Input(shape=(28, 28, 1)),
        keras.layers.Conv2D(64, (3, 3), strides=(2, 2), padding="same"),
        keras.layers.LeakyReLU(negative_slope=0.2),
        keras.layers.Conv2D(128, (3, 3), strides=(2, 2), padding="same"),
        keras.layers.LeakyReLU(negative_slope=0.2),
        keras.layers.GlobalMaxPooling2D(),
        keras.layers.Dense(1),
    ],
    name="discriminator",
)

# Example Generator Architecture
latent_dim = 128
generator = keras.Sequential(
    [
        keras.Input(shape=(latent_dim,)),
        keras.layers.Dense(7 * 7 * 128),
        keras.layers.LeakyReLU(negative_slope=0.2),
        keras.layers.Reshape((7, 7, 128)),
        keras.layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
        keras.layers.LeakyReLU(negative_slope=0.2),
        keras.layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding="same"),
        keras.layers.LeakyReLU(negative_slope=0.2),
        keras.layers.Conv2D(1, (7, 7), padding="same", activation="sigmoid"),
    ],
    name="generator",
)

To stop the entire hyperparameter search immediately when a critical condition is met (e.g., a bug is detected or a resource limit is hit), raise keras_tuner.errors.FatalError or one of its subclasses:

• FatalValueError
• FatalTypeError
• FatalRuntimeError

Raising these errors will terminate the search regardless of the max_consecutive_failed_trials setting. You should wrap your tuner.search() call in a try...except block to handle the termination gracefully.

def build_model(hp):
    # ... model definition ...
    num_params = model.count_params()
    if num_params > 1200:
        # This terminates the entire search process
        raise keras_tuner.errors.FatalError(
            f"Model too large! It contains {num_params} params."
        )
    return model

tuner = keras_tuner.GridSearch(
    hypermodel=build_model,
    objective="val_loss",
    max_retries_per_trial=3,
    max_consecutive_failed_trials=8,
)

try:
    tuner.search(x=x_data, y=y_data, ...)
except keras_tuner.errors.FatalError:
    print("The search is terminated.")

Because JAX cannot automatically differentiate through Pallas kernels, you must manually define the backward pass using @jax.custom_vjp.

To implement this:

1. Wrap your Pallas kernel with @jax.custom_vjp.
2. Define a forward function (fwd) that returns the output and a tuple of 'residuals' (the data needed for the backward pass, such as inputs or intermediate outputs).
3. Define a backward function (bwd) that accepts the residuals and the incoming gradient, calculates the gradients for each input, and returns them.
4. Register the functions using .defvjp(fwd, bwd).
5. Create a Keras layer that calls the wrapped trainable function in its call method.

# 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)