fractal_gpu.py

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def fractal_field(xi):
    xi_tensor = torch.tensor(xi, dtype=torch.float32, device=device)
    z = xi_tensor.clone()

    for _ in range(10):
        z = z**2 + xi_tensor

    return z.detach().cpu().numpy()