Practical Applications & Tutorials

Practical Applications & Tutorials#

This section provides detailed examples of using torch-secorder in various scenarios. For basic usage of foundational tools like Hessian-Vector Products (HVPs), Jacobian-Vector Products (JVPs), and basic curvature statistics (Hessian diagonal/trace), please refer to their respective API documentation in the Core and Approximations modules.

Analysis Examples#

Loss Landscape Visualization#

This example demonstrates how to compute and visualize the loss landscape of a neural network model. The loss landscape provides insights into the geometry of the optimization space and can help understand:

  • The smoothness and convexity of the loss function

  • The presence of local minima and saddle points

  • The difficulty of optimization in different regions

The example shows two types of visualizations:

  1. 1D Loss Surface: A slice of the loss landscape along a random direction, showing how the loss changes as we move away from the current parameter point.

  2. 2D Loss Surface: A contour plot showing the loss landscape in a 2D plane defined by two random directions, providing a more comprehensive view of the optimization space.

The visualization generates two files: - 1d_loss_surface.png: A line plot showing the loss along a single direction - 2d_loss_surface.png: A 3D surface plot showing the loss in a 2D plane

  1"""Example demonstrating loss landscape visualization.
  2
  3This script shows how to compute and visualize 1D slices and 2D contours
  4of the loss landscape for a simple model, using random directions.
  5
  6The example demonstrates:
  71. Computing 1D loss surface along a random direction
  82. Computing 2D loss surface using two random directions
  93. Visualizing both surfaces using matplotlib
 10
 11Requirements:
 12    - torch
 13    - matplotlib
 14    - numpy
 15"""
 16
 17import matplotlib.pyplot as plt
 18import torch
 19import torch.nn as nn
 20import torch.nn.functional as F
 21from matplotlib import cm
 22
 23from torch_secorder.analysis.landscape import (
 24    compute_loss_surface_1d,
 25    compute_loss_surface_2d,
 26    create_random_direction,
 27)
 28
 29
 30# 1. Define a Simple Model
 31class SimpleModel(nn.Module):
 32    def __init__(self):
 33        super().__init__()
 34        self.linear = nn.Linear(2, 1)
 35
 36    def forward(self, x):
 37        return self.linear(x)
 38
 39
 40# 2. Generate Synthetic Data
 41torch.manual_seed(42)
 42X_train = torch.randn(20, 2)
 43y_train = X_train @ torch.tensor([[2.0], [3.0]]) + 1.0 + torch.randn(20, 1) * 0.5
 44
 45
 46# 3. Instantiate Model and Loss Function
 47model = SimpleModel()
 48loss_fn = F.mse_loss
 49
 50
 51print("--- Loss Landscape Visualization Example ---")
 52
 53# 4. Compute 1D Loss Surface
 54print("\nComputing 1D loss surface...")
 55direction_1d = create_random_direction(model)
 56alphas_1d, losses_1d = compute_loss_surface_1d(
 57    model,
 58    loss_fn,
 59    X_train,
 60    y_train,
 61    direction_1d,
 62    alpha_range=(-2.0, 2.0),
 63    num_points=50,
 64)
 65
 66print("1D Loss Surface (first 5 points):\nAlphas: ", alphas_1d[:5].tolist())
 67print("Losses: ", losses_1d[:5].tolist())
 68
 69# Plot 1D Loss Surface
 70plt.figure(figsize=(8, 6))
 71plt.plot(alphas_1d.numpy(), losses_1d.numpy(), "b-", linewidth=2)
 72plt.xlabel("Alpha (Direction Scale)")
 73plt.ylabel("Loss")
 74plt.title("1D Loss Surface")
 75plt.grid(True)
 76plt.savefig("1d_loss_surface.png")
 77plt.close()
 78
 79
 80# 5. Compute 2D Loss Surface
 81print("\nComputing 2D loss surface...")
 82direction1_2d = create_random_direction(model)
 83direction2_2d = create_random_direction(model)
 84
 85# Ensure directions are not collinear (optional, but good for meaningful 2D surface)
 86# A simple way to get somewhat orthogonal directions: re-randomize if dot product is too high
 87# This is a heuristic, proper orthogonalization methods might be preferred for robustness
 88if (
 89    torch.dot(
 90        torch.cat([d.flatten() for d in direction1_2d]),
 91        torch.cat([d.flatten() for d in direction2_2d]),
 92    ).abs()
 93    > 0.5
 94):
 95    print("Adjusting second random direction for better orthogonality...")
 96    direction2_2d = create_random_direction(model)
 97
 98alphas_2d, betas_2d, losses_2d = compute_loss_surface_2d(
 99    model, loss_fn, X_train, y_train, direction1_2d, direction2_2d, num_points=25
100)
101
102print("2D Loss Surface (top-left 3x3 values):\n", losses_2d[:3, :3].tolist())
103
104# Plot 2D Loss Surface
105fig = plt.figure(figsize=(12, 10))
106ax = fig.add_subplot(111, projection="3d")
107A, B = torch.meshgrid(alphas_2d, betas_2d, indexing="ij")
108surf = ax.plot_surface(
109    A.numpy(),
110    B.numpy(),
111    losses_2d.numpy(),
112    cmap=cm.viridis,
113    edgecolor="none",
114    alpha=0.8,
115)
116ax.set_xlabel("Alpha (Direction 1)")
117ax.set_ylabel("Beta (Direction 2)")
118ax.set_zlabel("Loss")
119ax.set_title("2D Loss Surface")
120fig.colorbar(surf, shrink=0.5, aspect=5)
121plt.savefig("2d_loss_surface.png")
122plt.close()
123
124print(
125    "\nLoss landscape visualization complete. Check '1d_loss_surface.png' and '2d_loss_surface.png' for the plots."
126)

Per-Layer Curvature Analysis#

This example demonstrates how to extract per-layer Fisher diagonal elements, which form the basis for Kronecker-Factored Approximate Curvature (K-FAC) methods.

 1"""Example demonstrating per-layer curvature extraction and its relation to K-FAC.
 2
 3This script shows how to extract per-layer Fisher diagonal elements using per_layer_fisher_diagonal.
 4These per-layer blocks are the foundation for Kronecker-Factored Approximate Curvature (K-FAC).
 5"""
 6
 7import torch
 8import torch.nn as nn
 9import torch.nn.functional as F
10
11from torch_secorder.core.per_layer_curvature import per_layer_fisher_diagonal
12
13
14# Define a simple model with multiple layers
15class SimpleModel(nn.Module):
16    def __init__(self):
17        super().__init__()
18        self.linear1 = nn.Linear(10, 5)
19        self.linear2 = nn.Linear(5, 3)
20
21    def forward(self, x):
22        x = F.relu(self.linear1(x))
23        return self.linear2(x)
24
25
26# Generate synthetic data
27torch.manual_seed(42)
28X = torch.randn(8, 10)
29y = torch.randint(0, 3, (8,))  # 3 classes
30
31# Instantiate the model
32model = SimpleModel()
33
34# Compute per-layer Fisher diagonal
35# This extracts block-diagonal approximations of the Fisher Information Matrix
36# These blocks are the basis for K-FAC approximations
37layer_fisher_diagonals = per_layer_fisher_diagonal(model, F.cross_entropy, X, y)
38
39print("Per-Layer Fisher Diagonal (K-FAC basis):")
40for layer_name, diag in layer_fisher_diagonals.items():
41    print(
42        f"Layer: {layer_name}, Shape: {diag.shape}, First few values: {diag.flatten()[:5].tolist()}"
43    )
44
45print(
46    "\nK-FAC would further decompose these blocks into Kronecker products of smaller matrices."
47)

Optimization Examples#

Natural Gradient Descent#

This example demonstrates a simplified implementation of Natural Gradient Descent, leveraging the Fisher diagonal approximation from torch-secorder.

 1"""Example demonstrating Natural Gradient Descent using Fisher diagonal approximation."""
 2
 3import torch
 4import torch.nn as nn
 5import torch.optim as optim
 6
 7from torch_secorder.approximations.fisher_diagonal import empirical_fisher_diagonal
 8
 9
10# 1. Define a Simple Model
11class SimpleModel(nn.Module):
12    def __init__(self):
13        super().__init__()
14        self.linear = nn.Linear(2, 1)
15
16    def forward(self, x):
17        return self.linear(x)
18
19
20# 2. Generate Synthetic Data
21torch.manual_seed(42)
22X = torch.randn(100, 2)
23y = X @ torch.tensor([[1.0], [2.0]]) + 0.5 + torch.randn(100, 1) * 0.1
24
25# 3. Instantiate Model, Loss, and Optimizer
26model = SimpleModel()
27loss_fn = nn.MSELoss()
28
29# In a real NGD optimizer, this would be integrated directly.
30# Here, we demonstrate the manual step using Fisher diagonal.
31
32learning_rate = 0.01
33num_epochs = 100
34
35print("\n--- Training with Natural Gradient Descent (Approximation) ---")
36for epoch in range(num_epochs):
37    model.train()
38    optimizer = optim.SGD(
39        model.parameters(), lr=learning_rate
40    )  # We'll manually adjust gradients
41    optimizer.zero_grad()
42
43    # Forward pass
44    outputs = model(X)
45    loss = loss_fn(outputs, y)
46
47    # Backward pass to compute gradients
48    loss.backward(
49        create_graph=True
50    )  # create_graph=True is needed for higher-order derivatives like Fisher
51
52    # Compute Empirical Fisher Diagonal
53    # For NGD, we need F^{-1}g. Here, we approximate F as diagonal.
54    fisher_diagonal_list = empirical_fisher_diagonal(model, loss_fn, X, y)
55
56    # Manual Natural Gradient Update (approximated with diagonal Fisher)
57    with torch.no_grad():
58        for i, param in enumerate(model.parameters()):
59            if param.grad is not None:
60                # Natural gradient: F^{-1} * gradient
61                # If F is diagonal, F_ii = fisher_diagonal_list[i]_ii
62                # So F^{-1}_ii = 1 / fisher_diagonal_list[i]_ii
63                # Natural gradient update for each parameter:
64                # param_update = learning_rate * (param.grad / fisher_diagonal_list[i])
65                # param.data -= param_update
66                # A simpler diagonal update often used in practice for NGD is param.grad / (F_diag + epsilon)
67
68                # Let's use a simpler diagonal update: param.grad / (fisher_diag + epsilon)
69                epsilon = 1e-6  # For numerical stability
70                param.grad.data.div_(
71                    fisher_diagonal_list[i].clamp(min=epsilon)
72                )  # Element-wise division
73
74        optimizer.step()  # Apply the modified gradients
75
76    if (epoch + 1) % 10 == 0:
77        print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")
78
79print("\nTraining complete.")
80print(f"Final Loss: {loss.item():.4f}")
81print("Model parameters after NGD:")
82for name, param in model.named_parameters():
83    print(f"{name}: {param.data.squeeze()}")

Bayesian Neural Networks Inference#

This example illustrates how Fisher Information can be used in the context of Bayesian Neural Networks to approximate posterior variances of parameters.

  1"""Example demonstrating the use of Fisher Information for Bayesian Neural Networks (BNNs)."""
  2
  3import torch
  4import torch.nn as nn
  5import torch.nn.functional as F
  6
  7from torch_secorder.approximations import (
  8    empirical_fisher_diagonal,
  9    generalized_fisher_diagonal,
 10)
 11
 12
 13# 1. Define a Simple Bayesian Neural Network (for demonstration purposes)
 14# In a true BNN, weights would be distributions, but here we'll use point estimates
 15# and use Fisher to infer uncertainty or for approximate posterior.
 16class SimpleBNN(nn.Module):
 17    def __init__(self):
 18        super().__init__()
 19        self.linear1 = nn.Linear(10, 5)
 20        self.linear2 = nn.Linear(5, 1)
 21
 22    def forward(self, x):
 23        x = F.relu(self.linear1(x))
 24        return self.linear2(x)
 25
 26
 27# 2. Generate Synthetic Data
 28torch.manual_seed(42)
 29X_train = torch.randn(100, 10)
 30y_train = X_train @ torch.randn(10, 1) + 0.5 + torch.randn(100, 1) * 0.1
 31
 32X_test = torch.randn(20, 10)
 33
 34# 3. Instantiate Model and Loss Function
 35model = SimpleBNN()
 36
 37# For BNNs, often working with negative log-likelihood (NLL) directly
 38# For regression, this might be Gaussian NLL. For classification, CrossEntropy is NLL.
 39# Here, we'll use MSE for simplicity and demonstrate Fisher diagonal for it.
 40# Note: Generalized Fisher is typically for likelihoods (e.g., NLL in classification).
 41# For MSE in regression, Empirical Fisher is often more directly applicable.
 42
 43# Let's simulate a regression task with a Gaussian likelihood assumption
 44# where the loss is proportional to MSE (which is -log-likelihood for Gaussian with fixed variance)
 45# We'll adapt `generalized_fisher_diagonal` for a simplified regression NLL scenario.
 46
 47print("\n--- Estimating Uncertainty using Fisher Information (BNN Context) ---")
 48
 49# Train the model (simplified training loop)
 50optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
 51for epoch in range(50):
 52    optimizer.zero_grad()
 53    outputs = model(X_train)
 54    loss = F.mse_loss(outputs, y_train)
 55    loss.backward()
 56    optimizer.step()
 57    if (epoch + 1) % 10 == 0:
 58        print(f"Epoch [{epoch + 1}/50], Loss: {loss.item():.4f}")
 59
 60# 4. Use Fisher Information for Uncertainty Estimation (Approximation)
 61# For a BNN, the inverse Fisher (or Hessian) often approximates the posterior covariance.
 62# The diagonal of F^{-1} gives approximate posterior variances of parameters.
 63
 64print(
 65    "\n--- Computing Fisher Diagonal for Approximate Posterior Variance (Regression: MSE) ---"
 66)
 67
 68print(
 69    "\nUsing Empirical Fisher Diagonal for parameter uncertainty approximation (MSE regression)..."
 70)
 71try:
 72    # Using empirical_fisher_diagonal for a direct demonstration applicable to MSE
 73    fisher_diagonal_params = empirical_fisher_diagonal(
 74        model, F.mse_loss, X_train, y_train
 75    )
 76
 77    print("\nApproximate Posterior Variances (from inverse Fisher diagonal):")
 78    with torch.no_grad():
 79        for i, param in enumerate(model.parameters()):
 80            # Inverse of Fisher diagonal approximates variance of parameters
 81            # Ensure no division by zero if Fisher diagonal element is too small
 82            variance_approx = 1.0 / (
 83                fisher_diagonal_params[i] + 1e-8
 84            )  # Add epsilon for stability
 85            print(
 86                f"Parameter {i} variance approximation shape: {variance_approx.shape}"
 87            )
 88            print(f"  First few values: {variance_approx.flatten()[:5].tolist()}")
 89
 90except NotImplementedError as e:
 91    print(
 92        f"Error: {e}. Generalized Fisher diagonal for regression NLL might not be fully supported yet in `generalized_fisher_diagonal` for direct MSE loss."
 93    )
 94except Exception as e:
 95    print(f"An unexpected error occurred: {e}")
 96
 97print("\nBNN Fisher Information example complete.")
 98
 99# Add example for generalized_fisher_diagonal (classification use case)
100print("\n--- Using Generalized Fisher Diagonal (Classification Example) ---")
101
102
103# Create a simple classification model and data
104class SimpleClassifier(nn.Module):
105    def __init__(self):
106        super().__init__()
107        self.linear = nn.Linear(10, 3)
108
109    def forward(self, x):
110        return self.linear(x)
111
112
113clf_model = SimpleClassifier()
114clf_inputs = torch.randn(8, 10)
115clf_targets = torch.randint(0, 3, (8,))  # 3 classes
116clf_outputs = clf_model(clf_inputs)
117
118# Compute the Generalized Fisher Diagonal (NLL loss type)
119gfim_diagonal = generalized_fisher_diagonal(
120    clf_model, clf_outputs, clf_targets, loss_type="nll"
121)
122print("Generalized Fisher Diagonal (classification):")
123for i, diag in enumerate(gfim_diagonal):
124    print(f"  Param {i} shape: {diag.shape}, first few: {diag.flatten()[:5].tolist()}")

Gauss-Newton Optimization#

This example demonstrates a simplified optimization loop using the diagonal Gauss-Newton Matrix approximation.

 1"""Example demonstrating a simplified Gauss-Newton optimization step using GNM diagonal."""
 2
 3import torch
 4import torch.nn as nn
 5import torch.optim as optim
 6
 7from torch_secorder.approximations.gauss_newton import gauss_newton_matrix_approximation
 8
 9
10# 1. Define a Simple Model (e.g., for non-linear least squares)
11class SimpleNonLinearModel(nn.Module):
12    def __init__(self):
13        super().__init__()
14        # Let's make it slightly non-linear
15        self.linear1 = nn.Linear(2, 5)
16        self.relu = nn.ReLU()
17        self.linear2 = nn.Linear(5, 1)
18
19    def forward(self, x):
20        return self.linear2(self.relu(self.linear1(x)))
21
22
23# 2. Generate Synthetic Data
24torch.manual_seed(42)
25X = torch.randn(100, 2) * 5  # Input features
26y = (X[:, 0] ** 2 + X[:, 1] * 3 + 2 + torch.randn(100) * 0.5).unsqueeze(
27    1
28)  # Non-linear target
29
30# 3. Instantiate Model, Loss
31model = SimpleNonLinearModel()
32loss_fn = nn.MSELoss()  # Gauss-Newton is for least squares
33
34learning_rate = 0.01
35num_iterations = 100
36
37print("\n--- Training with Approximate Gauss-Newton (Diagonal GNM) ---")
38for iteration in range(num_iterations):
39    model.train()
40    optimizer = optim.SGD(
41        model.parameters(), lr=learning_rate
42    )  # We'll manually adjust gradients
43    optimizer.zero_grad()
44
45    # Forward pass
46    outputs = model(X)
47    loss = loss_fn(outputs, y)
48
49    # Backward pass for gradients
50    loss.backward(
51        create_graph=True
52    )  # create_graph=True if we want higher-order, not strictly needed for GNM diag
53
54    # Compute Diagonal of Gauss-Newton Matrix
55    # The inverse of the GNM is used in the update rule: delta_theta = - G^{-1} * gradient
56    # Here, we use diagonal GNM, so G^{-1}_ii = 1 / G_ii
57    gnm_diagonal_list = list(gauss_newton_matrix_approximation(model, loss_fn, X, y))
58
59    # Manual Gauss-Newton Update (approximated with diagonal GNM)
60    with torch.no_grad():
61        for i, param in enumerate(model.parameters()):
62            if param.grad is not None:
63                # Update: param.data -= learning_rate * (param.grad / (gnm_diag + epsilon))
64                epsilon = 1e-6  # For numerical stability
65                # Ensure gnm_diagonal_list[i] has the same shape as param.grad
66                param.grad.data.div_(
67                    gnm_diagonal_list[i].clamp(min=epsilon)
68                )  # Element-wise division
69
70        optimizer.step()  # Apply the modified gradients
71
72    if (iteration + 1) % 10 == 0:
73        print(f"Iteration [{iteration + 1}/{num_iterations}], Loss: {loss.item():.4f}")
74
75print("\nTraining complete.")
76print(f"Final Loss: {loss.item():.4f}")
77print("Model parameters after Gauss-Newton approximation:")
78for name, param in model.named_parameters():
79    print(f"{name}: {param.data.squeeze()}")

General Examples#

Quick Start#

A minimal example showing how to get started with torch-secorder.

 1"""Quick start example for torch-secorder.
 2
 3This example demonstrates basic usage of torch-secorder's core functionality:
 41. Computing Hessian-Vector Products (HVPs)
 52. Computing Jacobian-Vector Products (JVPs)
 63. Computing Hessian diagonal
 7"""
 8
 9import torch
10import torch.nn as nn
11import torch.nn.functional as F
12
13from torch_secorder.core import gvp, hessian_diagonal, hvp
14
15
16def main():
17    # 1. Create a simple model
18    model = nn.Sequential(nn.Linear(10, 5), nn.ReLU(), nn.Linear(5, 1))
19
20    # 2. Generate some random data
21    x = torch.randn(32, 10)
22    y = torch.randn(32, 1)
23
24    # 3. Compute loss
25    def get_loss():
26        return F.mse_loss(model(x), y)
27
28    # 4. Create a random direction vector
29    v = [torch.randn_like(p) for p in model.parameters()]
30
31    # 5. Compute HVP
32    hvp_result = hvp.model_hvp(model, F.mse_loss, x, y, v)
33    print("HVP computed successfully!")
34
35    # 6. Compute JVP
36    jvp_result = gvp.model_jvp(model, x, v)
37    print("JVP computed successfully!")
38
39    # 7. Compute Hessian diagonal
40    hessian_diag = hessian_diagonal(get_loss, list(model.parameters()))
41    print("Hessian diagonal computed successfully!")
42
43    # 8. Print shapes for verification
44    print("\nShapes of results:")
45    print(f"HVP: {[h.shape for h in hvp_result]}")
46    print(f"JVP: {[j.shape for j in jvp_result]}")
47    print(f"Hessian diagonal: {[h.shape for h in hessian_diag]}")
48
49
50if __name__ == "__main__":
51    main()

Custom Model Integration#

Demonstrates how to integrate torch-secorder with custom PyTorch models.

  1"""Example demonstrating integration with custom PyTorch models.
  2
  3This example shows how to use torch-secorder with custom PyTorch models,
  4including models with:
  51. Custom forward passes
  62. Multiple outputs
  73. Complex loss functions
  8"""
  9
 10import torch
 11import torch.nn as nn
 12import torch.nn.functional as F
 13
 14from torch_secorder.core import gvp, hessian_diagonal, hvp
 15
 16
 17class CustomModel(nn.Module):
 18    """A custom model with multiple outputs and a complex forward pass."""
 19
 20    def __init__(self):
 21        super().__init__()
 22        self.encoder = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 10))
 23        self.classifier = nn.Linear(10, 2)
 24        self.regressor = nn.Linear(10, 1)
 25
 26    def forward(self, x):
 27        features = self.encoder(x)
 28        return {
 29            "logits": self.classifier(features),
 30            "predictions": self.regressor(features),
 31        }
 32
 33
 34def custom_loss(outputs, targets):
 35    """A custom loss function combining classification and regression losses."""
 36    classification_loss = F.cross_entropy(outputs["logits"], targets["classes"])
 37    regression_loss = F.mse_loss(outputs["predictions"], targets["values"])
 38    return classification_loss + 0.1 * regression_loss
 39
 40
 41def main():
 42    # 1. Create the custom model
 43    model = CustomModel()
 44
 45    # 2. Generate random data
 46    x = torch.randn(32, 10)
 47    y = {"classes": torch.randint(0, 2, (32,)), "values": torch.randn(32, 1)}
 48
 49    # 3. Compute loss
 50    def get_loss():
 51        return custom_loss(model(x), y)
 52
 53    # 4. Create a random direction vector
 54    v = [torch.randn_like(p) for p in model.parameters()]
 55
 56    # 5. Compute HVP with custom model and loss
 57    hvp_result = hvp.model_hvp(model, custom_loss, x, y, v)
 58    print("HVP computed successfully with custom model!")
 59
 60    # 6. Compute JVP for each output separately
 61    # For models with multiple outputs, we need to compute JVP for each output tensor
 62    # We'll create wrapper models for each output
 63
 64    class LogitsModel(nn.Module):
 65        def __init__(self, base_model):
 66            super().__init__()
 67            self.base_model = base_model
 68
 69        def forward(self, x):
 70            return self.base_model(x)["logits"]
 71
 72    class PredictionsModel(nn.Module):
 73        def __init__(self, base_model):
 74            super().__init__()
 75            self.base_model = base_model
 76
 77        def forward(self, x):
 78            return self.base_model(x)["predictions"]
 79
 80    # Create wrapper models
 81    logits_model = LogitsModel(model)
 82    predictions_model = PredictionsModel(model)
 83
 84    # Compute JVPs using the wrapper models
 85    jvp_logits = gvp.model_jvp(logits_model, x, v)
 86    print("JVP computed successfully for logits!")
 87
 88    jvp_predictions = gvp.model_jvp(predictions_model, x, v)
 89    print("JVP computed successfully for predictions!")
 90
 91    # 7. Compute Hessian diagonal
 92    hessian_diag = hessian_diagonal(get_loss, list(model.parameters()))
 93    print("Hessian diagonal computed successfully with custom model!")
 94
 95    # 8. Print shapes for verification
 96    print("\nShapes of results:")
 97    print(f"HVP: {[h.shape for h in hvp_result]}")
 98    print(f"JVP (logits): {[j.shape for j in jvp_logits]}")
 99    print(f"JVP (predictions): {[j.shape for j in jvp_predictions]}")
100    print(f"Hessian diagonal: {[h.shape for h in hessian_diag]}")
101
102
103if __name__ == "__main__":
104    main()
Contents