warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/examples/tile/example_tile_convolution.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Tile Convolution
+#
+# Shows how to write a simple convolution kernel using Warp FFT tile
+# primitives.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+
+wp.set_module_options({"enable_backward": False})
+
+BLOCK_DIM = 64
+TILE_M = 1
+TILE_N = 128
+
+scale = wp.vec2d(wp.float64(1 / TILE_N), wp.float64(1 / TILE_N))
+
+
+@wp.func
+def filter(x: wp.vec2d):
+    return wp.cw_mul(x, scale)
+
+
+@wp.kernel
+def conv_tiled(x: wp.array2d(dtype=wp.vec2d), y: wp.array2d(dtype=wp.vec2d)):
+    i, j, _ = wp.tid()
+    a = wp.tile_load(x, shape=(TILE_M, TILE_N))
+    wp.tile_fft(a)
+    b = wp.tile_map(filter, a)
+    wp.tile_ifft(b)
+    wp.tile_store(y, b)
+
+
+if __name__ == "__main__":
+    wp.set_device("cuda:0")
+
+    rng = np.random.default_rng(42)
+
+    x_h = rng.standard_normal((TILE_M, TILE_N, 2), dtype=np.float64)
+    y_h = np.zeros_like(x_h)
+
+    x_wp = wp.array2d(x_h, dtype=wp.vec2d)
+    y_wp = wp.array2d(y_h, dtype=wp.vec2d)
+
+    wp.launch_tiled(conv_tiled, dim=[1, 1], inputs=[x_wp], outputs=[y_wp], block_dim=BLOCK_DIM)
+
+    # Since filter is 1/N, conv_tiled is a ~no-op
+    assert np.allclose(x_h, y_wp.numpy())

warp/examples/tile/example_tile_fft.py

@@ -0,0 +1,55 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Tile FFT
+#
+# Shows how to write a simple FFT kernel using Warp tile primitives.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+
+wp.set_module_options({"enable_backward": False})
+
+BLOCK_DIM = 8
+TILE_M = 1
+TILE_N = 32
+
+
+@wp.kernel
+def fft_tiled(x: wp.array2d(dtype=wp.vec2d), y: wp.array2d(dtype=wp.vec2d)):
+    i, j, _ = wp.tid()
+    a = wp.tile_load(x, shape=(TILE_M, TILE_N))
+    wp.tile_fft(a)
+    wp.tile_ifft(a)
+    wp.tile_store(y, a)
+
+
+if __name__ == "__main__":
+    wp.set_device("cuda:0")
+
+    x_h = np.ones((TILE_M, TILE_N, 2), dtype=np.float64)
+    x_h[:, :, 1] = 0
+    y_h = 3 * np.ones((TILE_M, TILE_N, 2), dtype=np.float64)
+    x_wp = wp.array2d(x_h, dtype=wp.vec2d)
+    y_wp = wp.array2d(y_h, dtype=wp.vec2d)
+
+    wp.launch_tiled(fft_tiled, dim=[1, 1], inputs=[x_wp], outputs=[y_wp], block_dim=BLOCK_DIM)
+
+    print("Inputs:\n", x_wp)  # [1+0i, 1+0i, 1+0i, ...]
+    print("Output:\n", y_wp)  # [32+0i, 0, 0, ...]

warp/examples/tile/example_tile_filtering.py

@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Tile Filtering
+#
+# Shows how to write a simple filtering kernel using Warp FFT tile
+# primitives.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+
+wp.set_module_options({"enable_backward": False})
+
+BLOCK_DIM = 128
+TILE_M = 1
+TILE_N = 512
+
+scale = wp.vec2d(wp.float64(1 / TILE_N), wp.float64(1 / TILE_N))
+
+
+def cplx(array):
+    return array[..., 0] + 1j * array[..., 1]
+
+
+@wp.func
+def cplx_prod(x: wp.vec2d, y: wp.vec2d):
+    return wp.cw_mul(wp.vec2d(x[0] * y[0] - x[1] * y[1], x[0] * y[1] + x[1] * y[0]), scale)
+
+
+@wp.kernel
+def conv_tiled(x: wp.array2d(dtype=wp.vec2d), y: wp.array2d(dtype=wp.vec2d), z: wp.array2d(dtype=wp.vec2d)):
+    i, j, _ = wp.tid()
+    a = wp.tile_load(x, shape=(TILE_M, TILE_N))
+    b = wp.tile_load(y, shape=(TILE_M, TILE_N))
+    wp.tile_fft(a)
+    c = wp.tile_map(cplx_prod, a, b)
+    wp.tile_ifft(c)
+    wp.tile_store(z, c)
+
+
+if __name__ == "__main__":
+    rng = np.random.default_rng(42)
+
+    # Create noisy input signal
+    t = np.linspace(0, 2 * np.pi, TILE_N, dtype=np.float64)
+    x = np.sin(t) + 0.5 * rng.random(TILE_N, dtype=np.float64)
+
+    # Create filter. This filter keeps only ~10% of the frequencies at the center
+    # of the spectrum.
+    f = np.ones_like(x)
+    freq = np.fft.fftfreq(TILE_N)
+    f[np.abs(freq) > 0.05] = 0.0
+    f[np.abs(freq) <= 0.05] = 1.0
+
+    # Create Warp input data
+    # We use vec2d to hold complex numbers
+    x_h = np.zeros((TILE_M, TILE_N, 2), dtype=np.float64)
+    f_h = np.zeros_like(x_h)
+    y_h = np.zeros_like(f_h)
+
+    x_h[:, :, 0] = x
+    f_h[:, :, 0] = f
+
+    x_wp = wp.array2d(x_h, dtype=wp.vec2d)
+    f_wp = wp.array2d(f_h, dtype=wp.vec2d)
+    y_wp = wp.array2d(y_h, dtype=wp.vec2d)
+
+    wp.launch_tiled(conv_tiled, dim=[1, 1], inputs=[x_wp, f_wp], outputs=[y_wp], block_dim=BLOCK_DIM)
+
+    # Extract output and compare with numpy
+    x_np = cplx(x_h)
+    f_np = cplx(f_h)
+    y_test = cplx(y_wp.numpy())
+    y_ref = np.fft.ifft(f_np * np.fft.fft(x_np))
+    assert np.allclose(y_ref, y_test)
+
+try:
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots(figsize=(10, 5))
+
+    ax.plot(
+        x,
+        color="#DDDDDD",
+        linewidth=2,
+        label="Original",
+    )
+    ax.plot(y_test[0, :].real, color="#76B900", linewidth=3, label="Smoothed")
+
+    ax.legend()
+    ax.grid(True)
+
+    plt.tight_layout()
+    plt.show()
+
+except ModuleNotFoundError:
+    print("Matplotlib not available; skipping figure")

warp/examples/tile/example_tile_matmul.py

@@ -0,0 +1,85 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Tile MatMul
+#
+# Shows how to write a simple GEMM kernel using Warp tile primitives.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+
+# tile size
+TILE_M = wp.constant(8)
+TILE_N = wp.constant(4)
+TILE_K = wp.constant(8)
+
+# num threads per-tile
+TILE_THREADS = 64
+
+
+@wp.kernel
+def tile_gemm(A: wp.array2d(dtype=wp.float32), B: wp.array2d(dtype=wp.float16), C: wp.array2d(dtype=wp.float64)):
+    # output tile index
+    i, j = wp.tid()
+
+    sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float64)
+
+    _M = A.shape[0]
+    _N = B.shape[1]
+    K = A.shape[1]
+
+    count = int(K / TILE_K)
+
+    for k in range(0, count):
+        a = wp.tile_load(A, shape=(TILE_M, TILE_K), offset=(i * TILE_M, k * TILE_K))
+        b = wp.tile_load(B, shape=(TILE_K, TILE_N), offset=(k * TILE_K, j * TILE_N))
+
+        # sum += a*b
+        wp.tile_matmul(a, b, sum)
+
+    wp.tile_store(C, sum, offset=(i * TILE_M, j * TILE_N))
+
+
+if __name__ == "__main__":
+    # generate some tile aligned matrix dimensions
+    M = TILE_M * 7
+    K = TILE_K * 6
+    N = TILE_N * 5
+
+    rng = np.random.default_rng(42)
+    A = rng.random((M, K), dtype=np.float32)
+    B = rng.random((K, N), dtype=np.float32).astype(np.float16)
+    C = np.zeros((M, N), dtype=np.float64)
+
+    A_wp = wp.array(A, requires_grad=True)
+    B_wp = wp.array(B, requires_grad=True)
+    C_wp = wp.array(C, requires_grad=True)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_gemm,
+            dim=(M // TILE_M, N // TILE_N),
+            inputs=[A_wp, B_wp],
+            outputs=[C_wp],
+            block_dim=TILE_THREADS,
+        )
+
+    assert np.allclose(C_wp.numpy(), A @ B, atol=1.0e-4)
+
+    print("Example matrix multiplication passed")

warp/examples/tile/example_tile_mlp.py

@@ -0,0 +1,383 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Image Multilayer Perceptron (MLP)
+#
+# Shows how to train a coordinate-based MLP on an image to predict the RGB
+# color at a given input position. By default, a positional encoding is
+# applied to the input coordinates to improve the ability of the MLP to
+# represent higher-frequency content. This can be disabled by passing the
+# '--no_encoding' option.
+#
+# References:
+#   Ben Mildenhall et al. 2021. NeRF: representing scenes
+#   as neural radiance fields for view synthesis. Commun. ACM 65, 1
+#   (January 2022), 99–106. https://doi.org/10.1145/3503250
+#
+###########################################################################
+
+import math
+import os
+
+import numpy as np
+from PIL import Image
+
+import warp as wp
+import warp.examples
+import warp.optim
+
+rng = np.random.default_rng(45)
+
+
+def create_layer(dim_in, dim_hid, dtype=float):
+    w = rng.uniform(-1.0 / np.sqrt(dim_in), 1.0 / np.sqrt(dim_in), (dim_hid, dim_in))
+    b = rng.uniform(-1.0 / np.sqrt(dim_in), 1.0 / np.sqrt(dim_in), (dim_hid, 1))
+
+    weights = wp.array(w, dtype=dtype, requires_grad=True)
+    bias = wp.array(b, dtype=dtype, requires_grad=True)
+
+    return (weights, bias)
+
+
+def create_array(dim_in, dim_hid, dtype=float):
+    s = rng.uniform(-1.0 / np.sqrt(dim_in), 1.0 / np.sqrt(dim_in), (dim_hid, dim_in))
+    a = wp.array(s, dtype=dtype, requires_grad=True)
+
+    return a
+
+
+# number of frequencies for the positional encoding
+NUM_FREQ = wp.constant(8)
+
+DIM_IN = wp.constant(4 * NUM_FREQ)  # sin,cos for both x,y at each frequenecy
+DIM_HID = 32
+DIM_OUT = 3
+
+# threads per-block
+NUM_THREADS = 32
+
+IMG_WIDTH = 512
+IMG_HEIGHT = 512
+
+BATCH_SIZE = min(1024, int((IMG_WIDTH * IMG_HEIGHT) / 8))
+
+# dtype for our weights and bias matrices
+dtype = wp.float16
+
+
+@wp.func
+def relu(x: dtype):
+    return wp.max(x, dtype(0.0))
+
+
+@wp.kernel
+def compute(
+    indices: wp.array(dtype=int),
+    weights_0: wp.array2d(dtype=dtype),
+    bias_0: wp.array2d(dtype=dtype),
+    weights_1: wp.array2d(dtype=dtype),
+    bias_1: wp.array2d(dtype=dtype),
+    weights_2: wp.array2d(dtype=dtype),
+    bias_2: wp.array2d(dtype=dtype),
+    weights_3: wp.array2d(dtype=dtype),
+    bias_3: wp.array2d(dtype=dtype),
+    reference: wp.array2d(dtype=float),
+    loss: wp.array1d(dtype=float),
+    out: wp.array2d(dtype=float),
+):
+    # batch indices
+    linear = indices[wp.tid()]
+
+    row = linear / IMG_WIDTH
+    col = linear % IMG_WIDTH
+
+    # normalize input coordinates to [-1, 1]
+    x = (float(row) / float(IMG_WIDTH) - 0.5) * 2.0
+    y = (float(col) / float(IMG_HEIGHT) - 0.5) * 2.0
+
+    local = wp.vector(dtype=dtype, length=DIM_IN)
+
+    # construct positional encoding
+    for s in range(NUM_FREQ):
+        scale = wp.pow(2.0, float(s)) * wp.pi
+
+        # x-coord
+        local[s * 4 + 0] = dtype(wp.sin(x * scale))
+        local[s * 4 + 1] = dtype(wp.cos(x * scale))
+        # y-coord
+        local[s * 4 + 2] = dtype(wp.sin(y * scale))
+        local[s * 4 + 3] = dtype(wp.cos(y * scale))
+
+    # tile feature vectors across the block, returns [dim(f), NUM_THREADS]
+    f = wp.tile(local)
+
+    # input layer
+    w0 = wp.tile_load(weights_0, shape=(DIM_HID, DIM_IN))
+    b0 = wp.tile_load(bias_0, shape=(DIM_HID, 1))
+    z = wp.tile_map(relu, wp.tile_matmul(w0, f) + wp.tile_broadcast(b0, shape=(DIM_HID, NUM_THREADS)))
+
+    # hidden layer
+    w1 = wp.tile_load(weights_1, shape=(DIM_HID, DIM_HID))
+    b1 = wp.tile_load(bias_1, shape=(DIM_HID, 1))
+    z = wp.tile_map(relu, wp.tile_matmul(w1, z) + wp.tile_broadcast(b1, shape=(DIM_HID, NUM_THREADS)))
+
+    w2 = wp.tile_load(weights_2, shape=(DIM_HID, DIM_HID))
+    b2 = wp.tile_load(bias_2, shape=(DIM_HID, 1))
+    z = wp.tile_map(relu, wp.tile_matmul(w2, z) + wp.tile_broadcast(b2, shape=(DIM_HID, NUM_THREADS)))
+
+    # output layer
+    w3 = wp.tile_load(weights_3, shape=(DIM_OUT, DIM_HID))
+    b3 = wp.tile_load(bias_3, shape=(DIM_OUT, 1))
+    o = wp.tile_map(relu, wp.tile_matmul(w3, z) + wp.tile_broadcast(b3, shape=(DIM_OUT, NUM_THREADS)))
+
+    # untile back to SIMT
+    output = wp.untile(o)
+
+    # compute error
+    error = wp.vec3(
+        float(output[0]) - reference[0, linear],
+        float(output[1]) - reference[1, linear],
+        float(output[2]) - reference[2, linear],
+    )
+
+    # write MSE loss
+    if loss:
+        wp.atomic_add(loss, 0, wp.length_sq(error) / float(3 * BATCH_SIZE))
+
+    #  write image output
+    if out:
+        for i in range(DIM_OUT):
+            out[i, linear] = float(output[i])
+
+
+class Example:
+    def __init__(self, train_iters):
+        self.weights_0, self.bias_0 = create_layer(DIM_IN, DIM_HID, dtype=dtype)
+        self.weights_1, self.bias_1 = create_layer(DIM_HID, DIM_HID, dtype=dtype)
+        self.weights_2, self.bias_2 = create_layer(DIM_HID, DIM_HID, dtype=dtype)
+        self.weights_3, self.bias_3 = create_layer(DIM_HID, DIM_OUT, dtype=dtype)
+
+        # reference
+        reference_path = os.path.join(wp.examples.get_asset_directory(), "pixel.jpg")
+        with Image.open(reference_path) as im:
+            reference_image = np.asarray(im.resize((IMG_WIDTH, IMG_HEIGHT)).convert("RGB")) / 255.0
+        self.reference = wp.array(reference_image.reshape(IMG_WIDTH * IMG_HEIGHT, 3).T, dtype=float)
+
+        # create randomized batch indices
+        indices = np.arange(0, IMG_WIDTH * IMG_HEIGHT, dtype=np.int32)
+        rng.shuffle(indices)
+        self.indices = wp.array(indices)
+
+        self.num_batches = int((IMG_WIDTH * IMG_HEIGHT) / BATCH_SIZE)
+        self.max_iters = train_iters
+        self.max_epochs = max(1, int(self.max_iters / self.num_batches))
+
+    def train_warp(self):
+        params = [
+            self.weights_0,
+            self.bias_0,
+            self.weights_1,
+            self.bias_1,
+            self.weights_2,
+            self.bias_2,
+            self.weights_3,
+            self.bias_3,
+        ]
+
+        optimizer_grads = [p.grad.flatten() for p in params]
+        optimizer_inputs = [p.flatten() for p in params]
+        optimizer = warp.optim.Adam(optimizer_inputs, lr=0.01)
+
+        loss = wp.zeros(1, dtype=float, requires_grad=True)
+        output = create_array(IMG_WIDTH * IMG_HEIGHT, DIM_OUT)
+
+        # capture graph for whole epoch
+        wp.capture_begin()
+
+        for b in range(0, IMG_WIDTH * IMG_HEIGHT, BATCH_SIZE):
+            loss.zero_()
+
+            with wp.Tape() as tape:
+                wp.launch(
+                    compute,
+                    dim=[BATCH_SIZE],
+                    inputs=[
+                        self.indices[b : b + BATCH_SIZE],
+                        self.weights_0,
+                        self.bias_0,
+                        self.weights_1,
+                        self.bias_1,
+                        self.weights_2,
+                        self.bias_2,
+                        self.weights_3,
+                        self.bias_3,
+                        self.reference,
+                        loss,
+                        None,
+                    ],
+                    block_dim=NUM_THREADS,
+                )
+
+            tape.backward(loss)
+            optimizer.step(optimizer_grads)
+            tape.zero()
+
+        graph = wp.capture_end()
+
+        with wp.ScopedTimer("Training"):
+            for i in range(self.max_epochs):
+                with wp.ScopedTimer("Epoch"):
+                    wp.capture_launch(graph)
+                    print(f"Epoch: {i} Loss: {loss.numpy()}")
+
+        # evaluate full image
+        wp.launch(
+            compute,
+            dim=[IMG_WIDTH * IMG_HEIGHT],
+            inputs=[
+                self.indices,
+                self.weights_0,
+                self.bias_0,
+                self.weights_1,
+                self.bias_1,
+                self.weights_2,
+                self.bias_2,
+                self.weights_3,
+                self.bias_3,
+                self.reference,
+                loss,
+                output,
+            ],
+            block_dim=NUM_THREADS,
+        )
+
+        self.save_image("example_tile_mlp.jpg", output.numpy())
+
+    def train_torch(self):
+        import torch as tc
+
+        weights_0 = tc.nn.Parameter(wp.to_torch(self.weights_0))
+        weights_1 = tc.nn.Parameter(wp.to_torch(self.weights_1))
+        weights_2 = tc.nn.Parameter(wp.to_torch(self.weights_2))
+        weights_3 = tc.nn.Parameter(wp.to_torch(self.weights_3))
+
+        bias_0 = tc.nn.Parameter(wp.to_torch(self.bias_0))
+        bias_1 = tc.nn.Parameter(wp.to_torch(self.bias_1))
+        bias_2 = tc.nn.Parameter(wp.to_torch(self.bias_2))
+        bias_3 = tc.nn.Parameter(wp.to_torch(self.bias_3))
+
+        indices = wp.to_torch(self.indices)
+        reference = wp.to_torch(self.reference)
+
+        optimizer = tc.optim.Adam(
+            [weights_0, bias_0, weights_1, bias_1, weights_2, bias_2, weights_3, bias_3],
+            capturable=True,
+            lr=0.0001,
+            betas=(0.9, 0.95),
+            eps=1.0e-6,
+        )
+
+        # generate frequency space encoding of pixels
+        # based on their linear index in the image
+        def encode(linear):
+            row = (linear // IMG_WIDTH).float()
+            col = (linear % IMG_WIDTH).float()
+
+            x = (row / float(IMG_WIDTH) - 0.5) * 2.0
+            y = (col / float(IMG_HEIGHT) - 0.5) * 2.0
+
+            encoding = tc.zeros((NUM_FREQ * 4, len(linear)), dtype=tc.float16, device="cuda")
+
+            for s in range(NUM_FREQ):
+                scale = math.pow(2.0, float(s)) * math.pi
+
+                # Directly write the computed values into the encoding tensor
+                encoding[s * 4 + 0, :] = tc.sin(scale * x)
+                encoding[s * 4 + 1, :] = tc.cos(scale * x)
+                encoding[s * 4 + 2, :] = tc.sin(scale * y)
+                encoding[s * 4 + 3, :] = tc.cos(scale * y)
+
+            return encoding
+
+        stream = tc.cuda.Stream()
+        graph = tc.cuda.CUDAGraph()
+
+        # warm-up
+        with tc.cuda.stream(stream):
+            f = tc.rand((NUM_FREQ * 4, BATCH_SIZE), dtype=tc.float16, device="cuda")
+            z = tc.relu(weights_0 @ f + bias_0)
+            z = tc.relu(weights_1 @ z + bias_1)
+            z = tc.relu(weights_2 @ z + bias_2)
+            z = tc.relu(weights_3 @ z + bias_3)
+            ref = tc.rand((3, BATCH_SIZE), dtype=tc.float16, device="cuda")
+            loss = tc.mean((z - ref) ** 2)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        with tc.cuda.graph(graph):
+            for b in range(0, IMG_WIDTH * IMG_HEIGHT, BATCH_SIZE):
+                linear = indices[b : b + BATCH_SIZE]
+
+                f = encode(linear)
+
+                z = tc.relu(weights_0 @ f + bias_0)
+                z = tc.relu(weights_1 @ z + bias_1)
+                z = tc.relu(weights_2 @ z + bias_2)
+                z = tc.relu(weights_3 @ z + bias_3)
+
+                ref = reference[:, linear]
+                loss = tc.mean((z - ref) ** 2)
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+        with wp.ScopedTimer("Training (Torch)"):
+            for _i in range(self.max_epochs):
+                with wp.ScopedTimer("Epoch"):
+                    graph.replay()
+
+                    print(loss)
+
+        f = encode(tc.arange(0, IMG_WIDTH * IMG_HEIGHT))
+        z = tc.relu(weights_0 @ f + bias_0)
+        z = tc.relu(weights_1 @ z + bias_1)
+        z = tc.relu(weights_2 @ z + bias_2)
+        z = tc.relu(weights_3 @ z + bias_3)
+
+        self.save_image("example_tile_mlp_torch.jpg", z.detach().cpu().numpy())
+
+    def save_image(self, name, output):
+        predicted_image = output.T.reshape(IMG_WIDTH, IMG_HEIGHT, 3)
+        predicted_image = (predicted_image * 255).astype(np.uint8)
+
+        predicted_image_pil = Image.fromarray(predicted_image)
+        predicted_image_pil.save(name)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--train_iters", type=int, default=20000, help="Total number of training iterations.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice("cuda:0"):
+        example = Example(args.train_iters)
+        example.train_warp()
+        # example.train_torch()