Value-and-Grads (GPU)

Open in Google Colab Download Notebook

Value-and-Grads (GPU)#

This notebook demonstrates automatic differentiation using Nabla, which enables efficient computation of gradients for optimization algorithms. The nb.vjp() function computes both the forward pass value and provides a function for the backward pass (Vector-Jacobian Product).

We try to target the GPU if available, but the code will also run on CPU.

Note: Nabla only works on GPU for functions that are wrapped with ``@nabla.jit``. This will trigger the MAX Graph compiler to create efficient GPU code. Eager GPU execution is not supported yet.

Setup and Imports#

[1]:

# Installation
import sys

IN_COLAB = "google.colab" in sys.modules

try:
    import nabla as nb
except ImportError:
    import subprocess

    subprocess.run(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "modular",
            "--extra-index-url",
            "https://download.pytorch.org/whl/cpu",
            "--index-url",
            "https://dl.modular.com/public/nightly/python/simple/",
        ],
        check=True,
    )
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "nabla-ml", "--upgrade"], check=True
    )
    import nabla as nb

print(
    f"🎉 Nabla is ready! Running on Python {sys.version_info.major}.{sys.version_info.minor}"
)

🎉 Nabla is ready! Running on Python 3.12

[3]:

# Device setup
device = nb.cpu() if nb.accelerator_count() == 0 else nb.accelerator()
print(f"Using device: {device}")

Using device: Device(type=gpu,id=0)

Define Function with Automatic Differentiation#

Create a JIT-compiled function that computes both the value and gradients.

[7]:

@nb.jit(show_graph=False)
def compute_with_gradients(x, y):
    def computation(x, y):
        return nb.sin(x * y) * 2

    value, vjp_fn = nb.vjp(computation, x, y)
    return value, vjp_fn(nb.ones_like(value))

Create Input Tensors#

Generate 2×3 tensors and move them to the target device.

[8]:

# Create tensors
a = nb.ndarange((2, 3)).to(device)
b = nb.ndarange((2, 3)).to(device)

Compute Values and Gradients#

Execute the function to get both the computed values and their gradients.

The output shows:

  • The first tensor contains the function values for sin(a * b) * 2.

  • The second tuple contains the gradients (∂f/∂a, ∂f/∂b) with respect to inputs a and b.

[9]:

# Compute and print results
value, grads = compute_with_gradients(a, b)
value, grads

[9]:

([[ 0.          1.6829418  -1.5136048 ]
  [ 0.82423794 -0.57580525 -0.2647073 ]]:f32[2,3],
 ([[ 0.         1.0806047 -2.6145742]
   [-5.4667797 -7.6612773  9.9120245]]:f32[2,3],
  [[ 0.         1.0806047 -2.6145742]
   [-5.4667797 -7.6612773  9.9120245]]:f32[2,3]))
Open in Google Colab Download Notebook

Note

💡 Want to run this yourself?

  • 🚀 Google Colab: No setup required, runs in your browser

  • 📥 Local Jupyter: Download and run with your own Python environment

Contents