Training a PINN on 2D PDE

In this tutorial we will go over using a PINN to solve 2D PDEs. We will be using the system from NeuralPDE Tutorials. However, we will be using our custom loss function and use nested AD capabilities of Lux.jl.

This is a demonstration of Lux.jl. For serious usecases of PINNs, please refer to the package: NeuralPDE.jl.

Package Imports

julia

using Lux,
    Optimisers,
    Random,
    Printf,
    Statistics,
    MLUtils,
    OnlineStats,
    CairoMakie,
    Reactant,
    Enzyme

const xdev = reactant_device(; force=true)
const cdev = cpu_device()

(::MLDataDevices.CPUDevice) (generic function with 1 method)

Problem Definition

Since Lux supports efficient nested AD upto 2nd order, we will rewrite the problem with first order derivatives, so that we can compute the gradients of the loss using 2nd order AD.

Define the Neural Networks

All the networks take 3 input variables and output a scalar value. Here, we will define a a wrapper over the 3 networks, so that we can train them using Training.TrainState.

julia

struct PINN{U,V,W} <: AbstractLuxContainerLayer{(:u, :v, :w)}
    u::U
    v::V
    w::W
end

function create_mlp(act, hidden_dims)
    return Chain(
        Dense(3 => hidden_dims, act),
        Dense(hidden_dims => hidden_dims, act),
        Dense(hidden_dims => hidden_dims, act),
        Dense(hidden_dims => 1),
    )
end

function PINN(; hidden_dims::Int=32)
    return PINN(
        create_mlp(tanh, hidden_dims),
        create_mlp(tanh, hidden_dims),
        create_mlp(tanh, hidden_dims),
    )
end

Main.var"##230".PINN

Define the Loss Functions

We will define a custom loss function to compute the loss using 2nd order AD. We will use the following loss function

julia

@views function physics_informed_loss_function(
    u::StatefulLuxLayer, v::StatefulLuxLayer, w::StatefulLuxLayer, xyt::AbstractArray
)
    ∂u_∂xyt = Enzyme.gradient(Enzyme.Reverse, sum ∘ u, xyt)[1]
    ∂u_∂x, ∂u_∂y, ∂u_∂t = ∂u_∂xyt[1:1, :], ∂u_∂xyt[2:2, :], ∂u_∂xyt[3:3, :]
    ∂v_∂x = Enzyme.gradient(Enzyme.Reverse, sum ∘ v, xyt)[1][1:1, :]
    v_xyt = v(xyt)
    ∂w_∂y = Enzyme.gradient(Enzyme.Reverse, sum ∘ w, xyt)[1][2:2, :]
    w_xyt = w(xyt)
    return (
        mean(abs2, ∂u_∂t .- ∂v_∂x .- ∂w_∂y) +
        mean(abs2, v_xyt .- ∂u_∂x) +
        mean(abs2, w_xyt .- ∂u_∂y)
    )
end

physics_informed_loss_function (generic function with 1 method)

Additionally, we need to compute the loss wrt the boundary conditions.

julia

function mse_loss_function(u::StatefulLuxLayer, target::AbstractArray, xyt::AbstractArray)
    return MSELoss()(u(xyt), target)
end

function loss_function(model, ps, st, (xyt, target_data, xyt_bc, target_bc))
    u_net = StatefulLuxLayer{true}(model.u, ps.u, st.u)
    v_net = StatefulLuxLayer{true}(model.v, ps.v, st.v)
    w_net = StatefulLuxLayer{true}(model.w, ps.w, st.w)
    physics_loss = physics_informed_loss_function(u_net, v_net, w_net, xyt)
    data_loss = mse_loss_function(u_net, target_data, xyt)
    bc_loss = mse_loss_function(u_net, target_bc, xyt_bc)
    loss = physics_loss + data_loss + bc_loss
    return (
        loss, (; u=u_net.st, v=v_net.st, w=w_net.st), (; physics_loss, data_loss, bc_loss)
    )
end

loss_function (generic function with 1 method)

Generate the Data

We will generate some random data to train the model on. We will take data on a square spatial and temporal domain $x \in [0, 2]$ , $y \in [0, 2]$ , and $t \in [0, 2]$ . Typically, you want to be smarter about the sampling process, but for the sake of simplicity, we will skip that.

julia

analytical_solution(x, y, t) = @. exp(x + y) * cos(x + y + 4t)
analytical_solution(xyt) = analytical_solution(xyt[1, :], xyt[2, :], xyt[3, :])

begin
    grid_len = 16

    grid = range(0.0f0, 2.0f0; length=grid_len)
    xyt = stack([[elem...] for elem in vec(collect(Iterators.product(grid, grid, grid)))])

    target_data = reshape(analytical_solution(xyt), 1, :)

    bc_len = 512

    x = collect(range(0.0f0, 2.0f0; length=bc_len))
    y = collect(range(0.0f0, 2.0f0; length=bc_len))
    t = collect(range(0.0f0, 2.0f0; length=bc_len))

    xyt_bc = hcat(
        stack((x, y, zeros(Float32, bc_len)); dims=1),
        stack((zeros(Float32, bc_len), y, t); dims=1),
        stack((ones(Float32, bc_len) .* 2, y, t); dims=1),
        stack((x, zeros(Float32, bc_len), t); dims=1),
        stack((x, ones(Float32, bc_len) .* 2, t); dims=1),
    )
    target_bc = reshape(analytical_solution(xyt_bc), 1, :)

    min_target_bc, max_target_bc = extrema(target_bc)
    min_data, max_data = extrema(target_data)
    min_pde_val, max_pde_val = min(min_data, min_target_bc), max(max_data, max_target_bc)

    xyt = (xyt .- minimum(xyt)) ./ (maximum(xyt) .- minimum(xyt))
    xyt_bc = (xyt_bc .- minimum(xyt_bc)) ./ (maximum(xyt_bc) .- minimum(xyt_bc))
    target_bc = (target_bc .- min_pde_val) ./ (max_pde_val - min_pde_val)
    target_data = (target_data .- min_pde_val) ./ (max_pde_val - min_pde_val)
end

Training

julia

function train_model(
    xyt,
    target_data,
    xyt_bc,
    target_bc;
    seed::Int=0,
    maxiters::Int=50000,
    hidden_dims::Int=32,
)
    rng = Random.default_rng()
    Random.seed!(rng, seed)

    pinn = PINN(; hidden_dims)
    ps, st = xdev(Lux.setup(rng, pinn))

    bc_dataloader = xdev(
        DataLoader((xyt_bc, target_bc); batchsize=32, shuffle=true, partial=false)
    )
    pde_dataloader = xdev(
        DataLoader((xyt, target_data); batchsize=32, shuffle=true, partial=false)
    )

    train_state = Training.TrainState(pinn, ps, st, Adam(0.05f0))
    lr = i -> i < 5000 ? 0.05f0 : (i < 10000 ? 0.005f0 : 0.0005f0)

    total_loss_tracker, physics_loss_tracker, data_loss_tracker, bc_loss_tracker = ntuple(
        _ -> OnlineStats.CircBuff(Float32, 32; rev=true), 4
    )

    iter = 1
    for ((xyt_batch, target_data_batch), (xyt_bc_batch, target_bc_batch)) in
        zip(Iterators.cycle(pde_dataloader), Iterators.cycle(bc_dataloader))
        Optimisers.adjust!(train_state, lr(iter))

        _, loss, stats, train_state = Training.single_train_step!(
            AutoEnzyme(),
            loss_function,
            (xyt_batch, target_data_batch, xyt_bc_batch, target_bc_batch),
            train_state,
        )

        fit!(total_loss_tracker, Float32(loss))
        fit!(physics_loss_tracker, Float32(stats.physics_loss))
        fit!(data_loss_tracker, Float32(stats.data_loss))
        fit!(bc_loss_tracker, Float32(stats.bc_loss))

        mean_loss = mean(OnlineStats.value(total_loss_tracker))
        mean_physics_loss = mean(OnlineStats.value(physics_loss_tracker))
        mean_data_loss = mean(OnlineStats.value(data_loss_tracker))
        mean_bc_loss = mean(OnlineStats.value(bc_loss_tracker))

        isnan(loss) && throw(ArgumentError("NaN Loss Detected"))

        if iter % 1000 == 1 || iter == maxiters
            @printf "Iteration: [%6d/%6d] \t Loss: %.9f (%.9f) \t Physics Loss: %.9f \
                     (%.9f) \t Data Loss: %.9f (%.9f) \t BC \
                     Loss: %.9f (%.9f)\n" iter maxiters loss mean_loss stats.physics_loss mean_physics_loss stats.data_loss mean_data_loss stats.bc_loss mean_bc_loss
        end

        iter += 1
        iter ≥ maxiters && break
    end

    return StatefulLuxLayer{true}(
        pinn, cdev(train_state.parameters), cdev(train_state.states)
    )
end

trained_model = train_model(xyt, target_data, xyt_bc, target_bc)
trained_u = Lux.testmode(
    StatefulLuxLayer{true}(trained_model.model.u, trained_model.ps.u, trained_model.st.u)
)

2025-03-28 04:44:33.167760: I external/xla/xla/service/service.cc:152] XLA service 0xb3b0180 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2025-03-28 04:44:33.167793: I external/xla/xla/service/service.cc:160]   StreamExecutor device (0): NVIDIA A100-PCIE-40GB MIG 1g.5gb, Compute Capability 8.0
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1743137073.168490 3481103 se_gpu_pjrt_client.cc:1039] Using BFC allocator.
I0000 00:00:1743137073.168557 3481103 gpu_helpers.cc:136] XLA backend allocating 3825205248 bytes on device 0 for BFCAllocator.
I0000 00:00:1743137073.168598 3481103 gpu_helpers.cc:177] XLA backend will use up to 1275068416 bytes on device 0 for CollectiveBFCAllocator.
I0000 00:00:1743137073.182124 3481103 cuda_dnn.cc:529] Loaded cuDNN version 90400
E0000 00:00:1743137370.266593 3481103 buffer_comparator.cc:156] Difference at 16: 0, expected 9.17386
E0000 00:00:1743137370.266655 3481103 buffer_comparator.cc:156] Difference at 17: 0, expected 8.63962
E0000 00:00:1743137370.266663 3481103 buffer_comparator.cc:156] Difference at 18: 0, expected 11.0889
E0000 00:00:1743137370.266671 3481103 buffer_comparator.cc:156] Difference at 19: 0, expected 8.96333
E0000 00:00:1743137370.266677 3481103 buffer_comparator.cc:156] Difference at 20: 0, expected 8.01067
E0000 00:00:1743137370.266684 3481103 buffer_comparator.cc:156] Difference at 21: 0, expected 11.0277
E0000 00:00:1743137370.266691 3481103 buffer_comparator.cc:156] Difference at 22: 0, expected 7.61027
E0000 00:00:1743137370.266697 3481103 buffer_comparator.cc:156] Difference at 23: 0, expected 7.70495
E0000 00:00:1743137370.266703 3481103 buffer_comparator.cc:156] Difference at 24: 0, expected 10.9025
E0000 00:00:1743137370.266710 3481103 buffer_comparator.cc:156] Difference at 25: 0, expected 9.9377
2025-03-28 04:49:30.266727: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.269465 3481103 buffer_comparator.cc:156] Difference at 16: 0, expected 9.17386
E0000 00:00:1743137370.269489 3481103 buffer_comparator.cc:156] Difference at 17: 0, expected 8.63962
E0000 00:00:1743137370.269496 3481103 buffer_comparator.cc:156] Difference at 18: 0, expected 11.0889
E0000 00:00:1743137370.269504 3481103 buffer_comparator.cc:156] Difference at 19: 0, expected 8.96333
E0000 00:00:1743137370.269510 3481103 buffer_comparator.cc:156] Difference at 20: 0, expected 8.01067
E0000 00:00:1743137370.269516 3481103 buffer_comparator.cc:156] Difference at 21: 0, expected 11.0277
E0000 00:00:1743137370.269523 3481103 buffer_comparator.cc:156] Difference at 22: 0, expected 7.61027
E0000 00:00:1743137370.269529 3481103 buffer_comparator.cc:156] Difference at 23: 0, expected 7.70495
E0000 00:00:1743137370.269535 3481103 buffer_comparator.cc:156] Difference at 24: 0, expected 10.9025
E0000 00:00:1743137370.269542 3481103 buffer_comparator.cc:156] Difference at 25: 0, expected 9.9377
2025-03-28 04:49:30.269552: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.271968 3481103 buffer_comparator.cc:156] Difference at 16: 0, expected 9.17386
E0000 00:00:1743137370.271984 3481103 buffer_comparator.cc:156] Difference at 17: 0, expected 8.63962
E0000 00:00:1743137370.271989 3481103 buffer_comparator.cc:156] Difference at 18: 0, expected 11.0889
E0000 00:00:1743137370.271994 3481103 buffer_comparator.cc:156] Difference at 19: 0, expected 8.96333
E0000 00:00:1743137370.271998 3481103 buffer_comparator.cc:156] Difference at 20: 0, expected 8.01067
E0000 00:00:1743137370.272002 3481103 buffer_comparator.cc:156] Difference at 21: 0, expected 11.0277
E0000 00:00:1743137370.272006 3481103 buffer_comparator.cc:156] Difference at 22: 0, expected 7.61027
E0000 00:00:1743137370.272011 3481103 buffer_comparator.cc:156] Difference at 23: 0, expected 7.70495
E0000 00:00:1743137370.272015 3481103 buffer_comparator.cc:156] Difference at 24: 0, expected 10.9025
E0000 00:00:1743137370.272021 3481103 buffer_comparator.cc:156] Difference at 25: 0, expected 9.9377
2025-03-28 04:49:30.272028: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.274243 3481103 buffer_comparator.cc:156] Difference at 32: 0, expected 9.26776
E0000 00:00:1743137370.274259 3481103 buffer_comparator.cc:156] Difference at 33: 0, expected 9.1446
E0000 00:00:1743137370.274263 3481103 buffer_comparator.cc:156] Difference at 34: 0, expected 9.79374
E0000 00:00:1743137370.274268 3481103 buffer_comparator.cc:156] Difference at 35: 0, expected 10.23
E0000 00:00:1743137370.274272 3481103 buffer_comparator.cc:156] Difference at 36: 0, expected 10.448
E0000 00:00:1743137370.274277 3481103 buffer_comparator.cc:156] Difference at 37: 0, expected 11.8017
E0000 00:00:1743137370.274281 3481103 buffer_comparator.cc:156] Difference at 38: 0, expected 9.69309
E0000 00:00:1743137370.274286 3481103 buffer_comparator.cc:156] Difference at 39: 0, expected 11.2277
E0000 00:00:1743137370.274290 3481103 buffer_comparator.cc:156] Difference at 40: 0, expected 10.0359
E0000 00:00:1743137370.274294 3481103 buffer_comparator.cc:156] Difference at 41: 0, expected 9.30719
2025-03-28 04:49:30.274301: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.276512 3481103 buffer_comparator.cc:156] Difference at 32: 0, expected 9.26776
E0000 00:00:1743137370.276528 3481103 buffer_comparator.cc:156] Difference at 33: 0, expected 9.1446
E0000 00:00:1743137370.276533 3481103 buffer_comparator.cc:156] Difference at 34: 0, expected 9.79374
E0000 00:00:1743137370.276537 3481103 buffer_comparator.cc:156] Difference at 35: 0, expected 10.23
E0000 00:00:1743137370.276542 3481103 buffer_comparator.cc:156] Difference at 36: 0, expected 10.448
E0000 00:00:1743137370.276546 3481103 buffer_comparator.cc:156] Difference at 37: 0, expected 11.8017
E0000 00:00:1743137370.276551 3481103 buffer_comparator.cc:156] Difference at 38: 0, expected 9.69309
E0000 00:00:1743137370.276555 3481103 buffer_comparator.cc:156] Difference at 39: 0, expected 11.2277
E0000 00:00:1743137370.276559 3481103 buffer_comparator.cc:156] Difference at 40: 0, expected 10.0359
E0000 00:00:1743137370.276564 3481103 buffer_comparator.cc:156] Difference at 41: 0, expected 9.30719
2025-03-28 04:49:30.276570: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.278780 3481103 buffer_comparator.cc:156] Difference at 32: 0, expected 9.26776
E0000 00:00:1743137370.278796 3481103 buffer_comparator.cc:156] Difference at 33: 0, expected 9.1446
E0000 00:00:1743137370.278800 3481103 buffer_comparator.cc:156] Difference at 34: 0, expected 9.79374
E0000 00:00:1743137370.278805 3481103 buffer_comparator.cc:156] Difference at 35: 0, expected 10.23
E0000 00:00:1743137370.278809 3481103 buffer_comparator.cc:156] Difference at 36: 0, expected 10.448
E0000 00:00:1743137370.278814 3481103 buffer_comparator.cc:156] Difference at 37: 0, expected 11.8017
E0000 00:00:1743137370.278818 3481103 buffer_comparator.cc:156] Difference at 38: 0, expected 9.69309
E0000 00:00:1743137370.278822 3481103 buffer_comparator.cc:156] Difference at 39: 0, expected 11.2277
E0000 00:00:1743137370.278827 3481103 buffer_comparator.cc:156] Difference at 40: 0, expected 10.0359
E0000 00:00:1743137370.278831 3481103 buffer_comparator.cc:156] Difference at 41: 0, expected 9.30719
2025-03-28 04:49:30.278838: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.296365 3481103 buffer_comparator.cc:156] Difference at 16: 0, expected 11.6059
E0000 00:00:1743137370.296399 3481103 buffer_comparator.cc:156] Difference at 17: 0, expected 14.502
E0000 00:00:1743137370.296403 3481103 buffer_comparator.cc:156] Difference at 18: 0, expected 11.2449
E0000 00:00:1743137370.296406 3481103 buffer_comparator.cc:156] Difference at 19: 0, expected 10.0998
E0000 00:00:1743137370.296409 3481103 buffer_comparator.cc:156] Difference at 20: 0, expected 14.0222
E0000 00:00:1743137370.296412 3481103 buffer_comparator.cc:156] Difference at 21: 0, expected 10.1321
E0000 00:00:1743137370.296415 3481103 buffer_comparator.cc:156] Difference at 22: 0, expected 10.2986
E0000 00:00:1743137370.296418 3481103 buffer_comparator.cc:156] Difference at 23: 0, expected 14.1109
E0000 00:00:1743137370.296421 3481103 buffer_comparator.cc:156] Difference at 24: 0, expected 13.3463
E0000 00:00:1743137370.296424 3481103 buffer_comparator.cc:156] Difference at 25: 0, expected 12.8369
2025-03-28 04:49:30.296432: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.299361 3481103 buffer_comparator.cc:156] Difference at 16: 0, expected 11.6059
E0000 00:00:1743137370.299373 3481103 buffer_comparator.cc:156] Difference at 17: 0, expected 14.502
E0000 00:00:1743137370.299376 3481103 buffer_comparator.cc:156] Difference at 18: 0, expected 11.2449
E0000 00:00:1743137370.299379 3481103 buffer_comparator.cc:156] Difference at 19: 0, expected 10.0998
E0000 00:00:1743137370.299382 3481103 buffer_comparator.cc:156] Difference at 20: 0, expected 14.0222
E0000 00:00:1743137370.299385 3481103 buffer_comparator.cc:156] Difference at 21: 0, expected 10.1321
E0000 00:00:1743137370.299388 3481103 buffer_comparator.cc:156] Difference at 22: 0, expected 10.2986
E0000 00:00:1743137370.299391 3481103 buffer_comparator.cc:156] Difference at 23: 0, expected 14.1109
E0000 00:00:1743137370.299394 3481103 buffer_comparator.cc:156] Difference at 24: 0, expected 13.3463
E0000 00:00:1743137370.299397 3481103 buffer_comparator.cc:156] Difference at 25: 0, expected 12.8369
2025-03-28 04:49:30.299401: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.302337 3481103 buffer_comparator.cc:156] Difference at 16: 0, expected 11.6059
E0000 00:00:1743137370.302349 3481103 buffer_comparator.cc:156] Difference at 17: 0, expected 14.502
E0000 00:00:1743137370.302352 3481103 buffer_comparator.cc:156] Difference at 18: 0, expected 11.2449
E0000 00:00:1743137370.302355 3481103 buffer_comparator.cc:156] Difference at 19: 0, expected 10.0998
E0000 00:00:1743137370.302358 3481103 buffer_comparator.cc:156] Difference at 20: 0, expected 14.0222
E0000 00:00:1743137370.302361 3481103 buffer_comparator.cc:156] Difference at 21: 0, expected 10.1321
E0000 00:00:1743137370.302364 3481103 buffer_comparator.cc:156] Difference at 22: 0, expected 10.2986
E0000 00:00:1743137370.302367 3481103 buffer_comparator.cc:156] Difference at 23: 0, expected 14.1109
E0000 00:00:1743137370.302370 3481103 buffer_comparator.cc:156] Difference at 24: 0, expected 13.3463
E0000 00:00:1743137370.302373 3481103 buffer_comparator.cc:156] Difference at 25: 0, expected 12.8369
2025-03-28 04:49:30.302377: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.305297 3481103 buffer_comparator.cc:156] Difference at 32: 0, expected 12.4
E0000 00:00:1743137370.305308 3481103 buffer_comparator.cc:156] Difference at 33: 0, expected 12.9454
E0000 00:00:1743137370.305312 3481103 buffer_comparator.cc:156] Difference at 34: 0, expected 12.9462
E0000 00:00:1743137370.305314 3481103 buffer_comparator.cc:156] Difference at 35: 0, expected 13.9775
E0000 00:00:1743137370.305317 3481103 buffer_comparator.cc:156] Difference at 36: 0, expected 15.0433
E0000 00:00:1743137370.305322 3481103 buffer_comparator.cc:156] Difference at 37: 0, expected 12.0589
E0000 00:00:1743137370.305325 3481103 buffer_comparator.cc:156] Difference at 38: 0, expected 14.4629
E0000 00:00:1743137370.305328 3481103 buffer_comparator.cc:156] Difference at 39: 0, expected 12.7671
E0000 00:00:1743137370.305331 3481103 buffer_comparator.cc:156] Difference at 40: 0, expected 12.3584
E0000 00:00:1743137370.305334 3481103 buffer_comparator.cc:156] Difference at 41: 0, expected 11.6002
2025-03-28 04:49:30.305338: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.308299 3481103 buffer_comparator.cc:156] Difference at 32: 0, expected 12.4
E0000 00:00:1743137370.308311 3481103 buffer_comparator.cc:156] Difference at 33: 0, expected 12.9454
E0000 00:00:1743137370.308314 3481103 buffer_comparator.cc:156] Difference at 34: 0, expected 12.9462
E0000 00:00:1743137370.308317 3481103 buffer_comparator.cc:156] Difference at 35: 0, expected 13.9775
E0000 00:00:1743137370.308320 3481103 buffer_comparator.cc:156] Difference at 36: 0, expected 15.0433
E0000 00:00:1743137370.308323 3481103 buffer_comparator.cc:156] Difference at 37: 0, expected 12.0589
E0000 00:00:1743137370.308326 3481103 buffer_comparator.cc:156] Difference at 38: 0, expected 14.4629
E0000 00:00:1743137370.308329 3481103 buffer_comparator.cc:156] Difference at 39: 0, expected 12.7671
E0000 00:00:1743137370.308332 3481103 buffer_comparator.cc:156] Difference at 40: 0, expected 12.3584
E0000 00:00:1743137370.308334 3481103 buffer_comparator.cc:156] Difference at 41: 0, expected 11.6002
2025-03-28 04:49:30.308339: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.311367 3481103 buffer_comparator.cc:156] Difference at 32: 0, expected 12.4
E0000 00:00:1743137370.311379 3481103 buffer_comparator.cc:156] Difference at 33: 0, expected 12.9454
E0000 00:00:1743137370.311382 3481103 buffer_comparator.cc:156] Difference at 34: 0, expected 12.9462
E0000 00:00:1743137370.311385 3481103 buffer_comparator.cc:156] Difference at 35: 0, expected 13.9775
E0000 00:00:1743137370.311388 3481103 buffer_comparator.cc:156] Difference at 36: 0, expected 15.0433
E0000 00:00:1743137370.311391 3481103 buffer_comparator.cc:156] Difference at 37: 0, expected 12.0589
E0000 00:00:1743137370.311394 3481103 buffer_comparator.cc:156] Difference at 38: 0, expected 14.4629
E0000 00:00:1743137370.311397 3481103 buffer_comparator.cc:156] Difference at 39: 0, expected 12.7671
E0000 00:00:1743137370.311400 3481103 buffer_comparator.cc:156] Difference at 40: 0, expected 12.3584
E0000 00:00:1743137370.311403 3481103 buffer_comparator.cc:156] Difference at 41: 0, expected 11.6002
2025-03-28 04:49:30.311407: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.325760 3481103 buffer_comparator.cc:156] Difference at 1024: 13.9198, expected 0.527629
E0000 00:00:1743137370.325779 3481103 buffer_comparator.cc:156] Difference at 1025: 12.3296, expected 0.544573
E0000 00:00:1743137370.325782 3481103 buffer_comparator.cc:156] Difference at 1026: 15.039, expected 0.31953
E0000 00:00:1743137370.325785 3481103 buffer_comparator.cc:156] Difference at 1027: 14.2357, expected 0.61527
E0000 00:00:1743137370.325788 3481103 buffer_comparator.cc:156] Difference at 1028: 10.1686, expected 1.06241
E0000 00:00:1743137370.325791 3481103 buffer_comparator.cc:156] Difference at 1029: 13.0036, expected 1.22093
E0000 00:00:1743137370.325795 3481103 buffer_comparator.cc:156] Difference at 1030: 13.3221, expected 0.553754
E0000 00:00:1743137370.325798 3481103 buffer_comparator.cc:156] Difference at 1031: 11.1666, expected 0.665092
E0000 00:00:1743137370.325801 3481103 buffer_comparator.cc:156] Difference at 1032: 14.2729, expected 1.13956
E0000 00:00:1743137370.325805 3481103 buffer_comparator.cc:156] Difference at 1033: 11.9378, expected 0.481583
2025-03-28 04:49:30.325811: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.333966 3481103 buffer_comparator.cc:156] Difference at 16: -nan, expected 7.32342
E0000 00:00:1743137370.333980 3481103 buffer_comparator.cc:156] Difference at 17: -nan, expected 10.2674
E0000 00:00:1743137370.333983 3481103 buffer_comparator.cc:156] Difference at 18: -nan, expected 8.76472
E0000 00:00:1743137370.333986 3481103 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.04411
E0000 00:00:1743137370.333989 3481103 buffer_comparator.cc:156] Difference at 20: -nan, expected 7.44869
E0000 00:00:1743137370.333992 3481103 buffer_comparator.cc:156] Difference at 21: -nan, expected 8.01524
E0000 00:00:1743137370.333994 3481103 buffer_comparator.cc:156] Difference at 22: -nan, expected 10.2221
E0000 00:00:1743137370.333997 3481103 buffer_comparator.cc:156] Difference at 23: -nan, expected 7.96118
E0000 00:00:1743137370.334000 3481103 buffer_comparator.cc:156] Difference at 24: -nan, expected 8.62185
E0000 00:00:1743137370.334003 3481103 buffer_comparator.cc:156] Difference at 25: -nan, expected 8.94114
2025-03-28 04:49:30.334008: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.336123 3481103 buffer_comparator.cc:156] Difference at 16: -nan, expected 7.32342
E0000 00:00:1743137370.336134 3481103 buffer_comparator.cc:156] Difference at 17: -nan, expected 10.2674
E0000 00:00:1743137370.336138 3481103 buffer_comparator.cc:156] Difference at 18: -nan, expected 8.76472
E0000 00:00:1743137370.336141 3481103 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.04411
E0000 00:00:1743137370.336143 3481103 buffer_comparator.cc:156] Difference at 20: -nan, expected 7.44869
E0000 00:00:1743137370.336146 3481103 buffer_comparator.cc:156] Difference at 21: -nan, expected 8.01524
E0000 00:00:1743137370.336149 3481103 buffer_comparator.cc:156] Difference at 22: -nan, expected 10.2221
E0000 00:00:1743137370.336152 3481103 buffer_comparator.cc:156] Difference at 23: -nan, expected 7.96118
E0000 00:00:1743137370.336154 3481103 buffer_comparator.cc:156] Difference at 24: -nan, expected 8.62185
E0000 00:00:1743137370.336157 3481103 buffer_comparator.cc:156] Difference at 25: -nan, expected 8.94114
2025-03-28 04:49:30.336161: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.338289 3481103 buffer_comparator.cc:156] Difference at 16: -nan, expected 7.32342
E0000 00:00:1743137370.338300 3481103 buffer_comparator.cc:156] Difference at 17: -nan, expected 10.2674
E0000 00:00:1743137370.338304 3481103 buffer_comparator.cc:156] Difference at 18: -nan, expected 8.76472
E0000 00:00:1743137370.338306 3481103 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.04411
E0000 00:00:1743137370.338309 3481103 buffer_comparator.cc:156] Difference at 20: -nan, expected 7.44869
E0000 00:00:1743137370.338312 3481103 buffer_comparator.cc:156] Difference at 21: -nan, expected 8.01524
E0000 00:00:1743137370.338314 3481103 buffer_comparator.cc:156] Difference at 22: -nan, expected 10.2221
E0000 00:00:1743137370.338317 3481103 buffer_comparator.cc:156] Difference at 23: -nan, expected 7.96118
E0000 00:00:1743137370.338320 3481103 buffer_comparator.cc:156] Difference at 24: -nan, expected 8.62185
E0000 00:00:1743137370.338322 3481103 buffer_comparator.cc:156] Difference at 25: -nan, expected 8.94114
2025-03-28 04:49:30.338327: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.340443 3481103 buffer_comparator.cc:156] Difference at 32: -nan, expected 7.79614
E0000 00:00:1743137370.340454 3481103 buffer_comparator.cc:156] Difference at 33: -nan, expected 8.40799
E0000 00:00:1743137370.340457 3481103 buffer_comparator.cc:156] Difference at 34: -nan, expected 7.73293
E0000 00:00:1743137370.340459 3481103 buffer_comparator.cc:156] Difference at 35: -nan, expected 9.69584
E0000 00:00:1743137370.340462 3481103 buffer_comparator.cc:156] Difference at 36: -nan, expected 8.80404
E0000 00:00:1743137370.340465 3481103 buffer_comparator.cc:156] Difference at 37: -nan, expected 10.1025
E0000 00:00:1743137370.340468 3481103 buffer_comparator.cc:156] Difference at 38: -nan, expected 8.25699
E0000 00:00:1743137370.340471 3481103 buffer_comparator.cc:156] Difference at 39: -nan, expected 8.00626
E0000 00:00:1743137370.340473 3481103 buffer_comparator.cc:156] Difference at 40: -nan, expected 8.86599
E0000 00:00:1743137370.340476 3481103 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.81407
2025-03-28 04:49:30.340480: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.342598 3481103 buffer_comparator.cc:156] Difference at 32: -nan, expected 7.79614
E0000 00:00:1743137370.342610 3481103 buffer_comparator.cc:156] Difference at 33: -nan, expected 8.40799
E0000 00:00:1743137370.342613 3481103 buffer_comparator.cc:156] Difference at 34: -nan, expected 7.73293
E0000 00:00:1743137370.342616 3481103 buffer_comparator.cc:156] Difference at 35: -nan, expected 9.69584
E0000 00:00:1743137370.342618 3481103 buffer_comparator.cc:156] Difference at 36: -nan, expected 8.80404
E0000 00:00:1743137370.342621 3481103 buffer_comparator.cc:156] Difference at 37: -nan, expected 10.1025
E0000 00:00:1743137370.342624 3481103 buffer_comparator.cc:156] Difference at 38: -nan, expected 8.25699
E0000 00:00:1743137370.342627 3481103 buffer_comparator.cc:156] Difference at 39: -nan, expected 8.00626
E0000 00:00:1743137370.342629 3481103 buffer_comparator.cc:156] Difference at 40: -nan, expected 8.86599
E0000 00:00:1743137370.342632 3481103 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.81407
2025-03-28 04:49:30.342636: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.344754 3481103 buffer_comparator.cc:156] Difference at 32: -nan, expected 7.79614
E0000 00:00:1743137370.344764 3481103 buffer_comparator.cc:156] Difference at 33: -nan, expected 8.40799
E0000 00:00:1743137370.344767 3481103 buffer_comparator.cc:156] Difference at 34: -nan, expected 7.73293
E0000 00:00:1743137370.344770 3481103 buffer_comparator.cc:156] Difference at 35: -nan, expected 9.69584
E0000 00:00:1743137370.344773 3481103 buffer_comparator.cc:156] Difference at 36: -nan, expected 8.80404
E0000 00:00:1743137370.344775 3481103 buffer_comparator.cc:156] Difference at 37: -nan, expected 10.1025
E0000 00:00:1743137370.344778 3481103 buffer_comparator.cc:156] Difference at 38: -nan, expected 8.25699
E0000 00:00:1743137370.344781 3481103 buffer_comparator.cc:156] Difference at 39: -nan, expected 8.00626
E0000 00:00:1743137370.344784 3481103 buffer_comparator.cc:156] Difference at 40: -nan, expected 8.86599
E0000 00:00:1743137370.344786 3481103 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.81407
2025-03-28 04:49:30.344791: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.372851 3481103 buffer_comparator.cc:156] Difference at 16: -nan, expected 9.278
E0000 00:00:1743137370.372895 3481103 buffer_comparator.cc:156] Difference at 17: -nan, expected 9.69724
E0000 00:00:1743137370.372900 3481103 buffer_comparator.cc:156] Difference at 18: -nan, expected 10.7774
E0000 00:00:1743137370.372903 3481103 buffer_comparator.cc:156] Difference at 19: -nan, expected 10.3636
E0000 00:00:1743137370.372906 3481103 buffer_comparator.cc:156] Difference at 20: -nan, expected 7.56533
E0000 00:00:1743137370.372909 3481103 buffer_comparator.cc:156] Difference at 21: -nan, expected 8.17756
E0000 00:00:1743137370.372911 3481103 buffer_comparator.cc:156] Difference at 22: -nan, expected 11.0551
E0000 00:00:1743137370.372914 3481103 buffer_comparator.cc:156] Difference at 23: -nan, expected 11.628
E0000 00:00:1743137370.372917 3481103 buffer_comparator.cc:156] Difference at 24: -nan, expected 10.6715
E0000 00:00:1743137370.372920 3481103 buffer_comparator.cc:156] Difference at 25: -nan, expected 10.8131
2025-03-28 04:49:30.372929: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.375058 3481103 buffer_comparator.cc:156] Difference at 16: -nan, expected 9.278
E0000 00:00:1743137370.375069 3481103 buffer_comparator.cc:156] Difference at 17: -nan, expected 9.69724
E0000 00:00:1743137370.375072 3481103 buffer_comparator.cc:156] Difference at 18: -nan, expected 10.7774
E0000 00:00:1743137370.375075 3481103 buffer_comparator.cc:156] Difference at 19: -nan, expected 10.3636
E0000 00:00:1743137370.375078 3481103 buffer_comparator.cc:156] Difference at 20: -nan, expected 7.56533
E0000 00:00:1743137370.375081 3481103 buffer_comparator.cc:156] Difference at 21: -nan, expected 8.17756
E0000 00:00:1743137370.375084 3481103 buffer_comparator.cc:156] Difference at 22: -nan, expected 11.0551
E0000 00:00:1743137370.375086 3481103 buffer_comparator.cc:156] Difference at 23: -nan, expected 11.628
E0000 00:00:1743137370.375089 3481103 buffer_comparator.cc:156] Difference at 24: -nan, expected 10.6715
E0000 00:00:1743137370.375092 3481103 buffer_comparator.cc:156] Difference at 25: -nan, expected 10.8131
2025-03-28 04:49:30.375096: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.377221 3481103 buffer_comparator.cc:156] Difference at 16: -nan, expected 9.278
E0000 00:00:1743137370.377232 3481103 buffer_comparator.cc:156] Difference at 17: -nan, expected 9.69724
E0000 00:00:1743137370.377235 3481103 buffer_comparator.cc:156] Difference at 18: -nan, expected 10.7774
E0000 00:00:1743137370.377238 3481103 buffer_comparator.cc:156] Difference at 19: -nan, expected 10.3636
E0000 00:00:1743137370.377241 3481103 buffer_comparator.cc:156] Difference at 20: -nan, expected 7.56533
E0000 00:00:1743137370.377244 3481103 buffer_comparator.cc:156] Difference at 21: -nan, expected 8.17756
E0000 00:00:1743137370.377246 3481103 buffer_comparator.cc:156] Difference at 22: -nan, expected 11.0551
E0000 00:00:1743137370.377249 3481103 buffer_comparator.cc:156] Difference at 23: -nan, expected 11.628
E0000 00:00:1743137370.377252 3481103 buffer_comparator.cc:156] Difference at 24: -nan, expected 10.6715
E0000 00:00:1743137370.377255 3481103 buffer_comparator.cc:156] Difference at 25: -nan, expected 10.8131
2025-03-28 04:49:30.377259: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.379373 3481103 buffer_comparator.cc:156] Difference at 32: -nan, expected 10.5897
E0000 00:00:1743137370.379384 3481103 buffer_comparator.cc:156] Difference at 33: -nan, expected 8.79857
E0000 00:00:1743137370.379387 3481103 buffer_comparator.cc:156] Difference at 34: -nan, expected 10.7009
E0000 00:00:1743137370.379390 3481103 buffer_comparator.cc:156] Difference at 35: -nan, expected 10.0423
E0000 00:00:1743137370.379393 3481103 buffer_comparator.cc:156] Difference at 36: -nan, expected 10.3055
E0000 00:00:1743137370.379398 3481103 buffer_comparator.cc:156] Difference at 37: -nan, expected 11.1027
E0000 00:00:1743137370.379401 3481103 buffer_comparator.cc:156] Difference at 38: -nan, expected 9.64939
E0000 00:00:1743137370.379403 3481103 buffer_comparator.cc:156] Difference at 39: -nan, expected 10.2558
E0000 00:00:1743137370.379406 3481103 buffer_comparator.cc:156] Difference at 40: -nan, expected 8.38055
E0000 00:00:1743137370.379409 3481103 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.92221
2025-03-28 04:49:30.379413: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.381554 3481103 buffer_comparator.cc:156] Difference at 32: -nan, expected 10.5897
E0000 00:00:1743137370.381567 3481103 buffer_comparator.cc:156] Difference at 33: -nan, expected 8.79857
E0000 00:00:1743137370.381570 3481103 buffer_comparator.cc:156] Difference at 34: -nan, expected 10.7009
E0000 00:00:1743137370.381573 3481103 buffer_comparator.cc:156] Difference at 35: -nan, expected 10.0423
E0000 00:00:1743137370.381576 3481103 buffer_comparator.cc:156] Difference at 36: -nan, expected 10.3055
E0000 00:00:1743137370.381579 3481103 buffer_comparator.cc:156] Difference at 37: -nan, expected 11.1027
E0000 00:00:1743137370.381581 3481103 buffer_comparator.cc:156] Difference at 38: -nan, expected 9.64939
E0000 00:00:1743137370.381584 3481103 buffer_comparator.cc:156] Difference at 39: -nan, expected 10.2558
E0000 00:00:1743137370.381587 3481103 buffer_comparator.cc:156] Difference at 40: -nan, expected 8.38055
E0000 00:00:1743137370.381589 3481103 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.92221
2025-03-28 04:49:30.381594: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.383741 3481103 buffer_comparator.cc:156] Difference at 32: -nan, expected 10.5897
E0000 00:00:1743137370.383753 3481103 buffer_comparator.cc:156] Difference at 33: -nan, expected 8.79857
E0000 00:00:1743137370.383757 3481103 buffer_comparator.cc:156] Difference at 34: -nan, expected 10.7009
E0000 00:00:1743137370.383760 3481103 buffer_comparator.cc:156] Difference at 35: -nan, expected 10.0423
E0000 00:00:1743137370.383762 3481103 buffer_comparator.cc:156] Difference at 36: -nan, expected 10.3055
E0000 00:00:1743137370.383765 3481103 buffer_comparator.cc:156] Difference at 37: -nan, expected 11.1027
E0000 00:00:1743137370.383768 3481103 buffer_comparator.cc:156] Difference at 38: -nan, expected 9.64939
E0000 00:00:1743137370.383771 3481103 buffer_comparator.cc:156] Difference at 39: -nan, expected 10.2558
E0000 00:00:1743137370.383773 3481103 buffer_comparator.cc:156] Difference at 40: -nan, expected 8.38055
E0000 00:00:1743137370.383776 3481103 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.92221
2025-03-28 04:49:30.383781: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.385909 3481103 buffer_comparator.cc:156] Difference at 64: -nan, expected 8.4725
E0000 00:00:1743137370.385920 3481103 buffer_comparator.cc:156] Difference at 65: -nan, expected 8.00729
E0000 00:00:1743137370.385923 3481103 buffer_comparator.cc:156] Difference at 66: -nan, expected 10.5558
E0000 00:00:1743137370.385926 3481103 buffer_comparator.cc:156] Difference at 67: -nan, expected 7.97013
E0000 00:00:1743137370.385928 3481103 buffer_comparator.cc:156] Difference at 68: -nan, expected 9.17677
E0000 00:00:1743137370.385931 3481103 buffer_comparator.cc:156] Difference at 69: -nan, expected 6.91094
E0000 00:00:1743137370.385934 3481103 buffer_comparator.cc:156] Difference at 70: -nan, expected 7.85311
E0000 00:00:1743137370.385936 3481103 buffer_comparator.cc:156] Difference at 71: -nan, expected 8.69839
E0000 00:00:1743137370.385941 3481103 buffer_comparator.cc:156] Difference at 72: -nan, expected 9.37233
E0000 00:00:1743137370.385944 3481103 buffer_comparator.cc:156] Difference at 73: -nan, expected 7.56939
2025-03-28 04:49:30.385948: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.388073 3481103 buffer_comparator.cc:156] Difference at 64: -nan, expected 8.4725
E0000 00:00:1743137370.388084 3481103 buffer_comparator.cc:156] Difference at 65: -nan, expected 8.00729
E0000 00:00:1743137370.388087 3481103 buffer_comparator.cc:156] Difference at 66: -nan, expected 10.5558
E0000 00:00:1743137370.388090 3481103 buffer_comparator.cc:156] Difference at 67: -nan, expected 7.97013
E0000 00:00:1743137370.388092 3481103 buffer_comparator.cc:156] Difference at 68: -nan, expected 9.17677
E0000 00:00:1743137370.388095 3481103 buffer_comparator.cc:156] Difference at 69: -nan, expected 6.91094
E0000 00:00:1743137370.388098 3481103 buffer_comparator.cc:156] Difference at 70: -nan, expected 7.85311
E0000 00:00:1743137370.388100 3481103 buffer_comparator.cc:156] Difference at 71: -nan, expected 8.69839
E0000 00:00:1743137370.388103 3481103 buffer_comparator.cc:156] Difference at 72: -nan, expected 9.37233
E0000 00:00:1743137370.388106 3481103 buffer_comparator.cc:156] Difference at 73: -nan, expected 7.56939
2025-03-28 04:49:30.388110: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.390241 3481103 buffer_comparator.cc:156] Difference at 64: -nan, expected 8.4725
E0000 00:00:1743137370.390252 3481103 buffer_comparator.cc:156] Difference at 65: -nan, expected 8.00729
E0000 00:00:1743137370.390256 3481103 buffer_comparator.cc:156] Difference at 66: -nan, expected 10.5558
E0000 00:00:1743137370.390259 3481103 buffer_comparator.cc:156] Difference at 67: -nan, expected 7.97013
E0000 00:00:1743137370.390263 3481103 buffer_comparator.cc:156] Difference at 68: -nan, expected 9.17677
E0000 00:00:1743137370.390266 3481103 buffer_comparator.cc:156] Difference at 69: -nan, expected 6.91094
E0000 00:00:1743137370.390268 3481103 buffer_comparator.cc:156] Difference at 70: -nan, expected 7.85311
E0000 00:00:1743137370.390271 3481103 buffer_comparator.cc:156] Difference at 71: -nan, expected 8.69839
E0000 00:00:1743137370.390274 3481103 buffer_comparator.cc:156] Difference at 72: -nan, expected 9.37233
E0000 00:00:1743137370.390278 3481103 buffer_comparator.cc:156] Difference at 73: -nan, expected 7.56939
2025-03-28 04:49:30.390283: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.397292 3481103 buffer_comparator.cc:156] Difference at 17: 9.32088, expected 7.11199
E0000 00:00:1743137370.397328 3481103 buffer_comparator.cc:156] Difference at 21: 7.08398, expected 9.39119
E0000 00:00:1743137370.397331 3481103 buffer_comparator.cc:156] Difference at 22: 7.37219, expected 9.38334
E0000 00:00:1743137370.397334 3481103 buffer_comparator.cc:156] Difference at 24: 6.74966, expected 8.44799
E0000 00:00:1743137370.397337 3481103 buffer_comparator.cc:156] Difference at 26: 7.50065, expected 9.26022
E0000 00:00:1743137370.397340 3481103 buffer_comparator.cc:156] Difference at 27: 7.63766, expected 9.09301
E0000 00:00:1743137370.397343 3481103 buffer_comparator.cc:156] Difference at 28: 7.32136, expected 9.93336
E0000 00:00:1743137370.397346 3481103 buffer_comparator.cc:156] Difference at 29: 6.62533, expected 9.39435
E0000 00:00:1743137370.397348 3481103 buffer_comparator.cc:156] Difference at 30: 8.2956, expected 9.51389
E0000 00:00:1743137370.397351 3481103 buffer_comparator.cc:156] Difference at 34: 10.3938, expected 7.48572
2025-03-28 04:49:30.397360: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.399498 3481103 buffer_comparator.cc:156] Difference at 17: 9.32088, expected 7.11199
E0000 00:00:1743137370.399514 3481103 buffer_comparator.cc:156] Difference at 21: 7.08398, expected 9.39119
E0000 00:00:1743137370.399518 3481103 buffer_comparator.cc:156] Difference at 22: 7.37219, expected 9.38334
E0000 00:00:1743137370.399520 3481103 buffer_comparator.cc:156] Difference at 24: 6.74966, expected 8.44799
E0000 00:00:1743137370.399523 3481103 buffer_comparator.cc:156] Difference at 26: 7.50065, expected 9.26022
E0000 00:00:1743137370.399526 3481103 buffer_comparator.cc:156] Difference at 27: 7.63766, expected 9.09301
E0000 00:00:1743137370.399529 3481103 buffer_comparator.cc:156] Difference at 28: 7.32136, expected 9.93336
E0000 00:00:1743137370.399532 3481103 buffer_comparator.cc:156] Difference at 29: 6.62533, expected 9.39435
E0000 00:00:1743137370.399534 3481103 buffer_comparator.cc:156] Difference at 30: 8.2956, expected 9.51389
E0000 00:00:1743137370.399537 3481103 buffer_comparator.cc:156] Difference at 34: 10.3938, expected 7.48572
2025-03-28 04:49:30.399544: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.401671 3481103 buffer_comparator.cc:156] Difference at 34: 10.3938, expected 7.48572
E0000 00:00:1743137370.401683 3481103 buffer_comparator.cc:156] Difference at 39: 8.70204, expected 7.6147
E0000 00:00:1743137370.401686 3481103 buffer_comparator.cc:156] Difference at 40: 10.1065, expected 6.04756
E0000 00:00:1743137370.401689 3481103 buffer_comparator.cc:156] Difference at 41: 10.7572, expected 7.8893
E0000 00:00:1743137370.401692 3481103 buffer_comparator.cc:156] Difference at 43: 10.7753, expected 7.69612
E0000 00:00:1743137370.401695 3481103 buffer_comparator.cc:156] Difference at 44: 9.91027, expected 8.503
E0000 00:00:1743137370.401698 3481103 buffer_comparator.cc:156] Difference at 45: 10.2077, expected 8.12358
E0000 00:00:1743137370.401701 3481103 buffer_comparator.cc:156] Difference at 46: 10.7749, expected 7.84264
E0000 00:00:1743137370.401704 3481103 buffer_comparator.cc:156] Difference at 48: 9.44041, expected 7.26336
E0000 00:00:1743137370.401707 3481103 buffer_comparator.cc:156] Difference at 49: 9.96615, expected 6.98996
2025-03-28 04:49:30.401711: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.403857 3481103 buffer_comparator.cc:156] Difference at 34: 10.3938, expected 7.48572
E0000 00:00:1743137370.403868 3481103 buffer_comparator.cc:156] Difference at 39: 8.70204, expected 7.6147
E0000 00:00:1743137370.403871 3481103 buffer_comparator.cc:156] Difference at 40: 10.1065, expected 6.04756
E0000 00:00:1743137370.403875 3481103 buffer_comparator.cc:156] Difference at 41: 10.7572, expected 7.8893
E0000 00:00:1743137370.403878 3481103 buffer_comparator.cc:156] Difference at 43: 10.7753, expected 7.69612
E0000 00:00:1743137370.403881 3481103 buffer_comparator.cc:156] Difference at 44: 9.91027, expected 8.503
E0000 00:00:1743137370.403883 3481103 buffer_comparator.cc:156] Difference at 45: 10.2077, expected 8.12358
E0000 00:00:1743137370.403886 3481103 buffer_comparator.cc:156] Difference at 46: 10.7749, expected 7.84264
E0000 00:00:1743137370.403889 3481103 buffer_comparator.cc:156] Difference at 48: 9.44041, expected 7.26336
E0000 00:00:1743137370.403892 3481103 buffer_comparator.cc:156] Difference at 49: 9.96615, expected 6.98996
2025-03-28 04:49:30.403896: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.406034 3481103 buffer_comparator.cc:156] Difference at 67: 10.5021, expected 7.70695
E0000 00:00:1743137370.406047 3481103 buffer_comparator.cc:156] Difference at 68: 10.1174, expected 9.00042
E0000 00:00:1743137370.406062 3481103 buffer_comparator.cc:156] Difference at 69: 9.89636, expected 8.345
E0000 00:00:1743137370.406065 3481103 buffer_comparator.cc:156] Difference at 70: 10.5296, expected 8.52479
E0000 00:00:1743137370.406068 3481103 buffer_comparator.cc:156] Difference at 71: 8.1412, expected 7.08541
E0000 00:00:1743137370.406071 3481103 buffer_comparator.cc:156] Difference at 72: 6.85179, expected 8.47292
E0000 00:00:1743137370.406074 3481103 buffer_comparator.cc:156] Difference at 77: 7.45304, expected 9.28323
E0000 00:00:1743137370.406077 3481103 buffer_comparator.cc:156] Difference at 80: 7.39802, expected 9.4894
E0000 00:00:1743137370.406080 3481103 buffer_comparator.cc:156] Difference at 81: 9.68936, expected 7.49939
E0000 00:00:1743137370.406082 3481103 buffer_comparator.cc:156] Difference at 82: 9.2677, expected 6.28629
2025-03-28 04:49:30.406087: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.408207 3481103 buffer_comparator.cc:156] Difference at 67: 10.5021, expected 7.70695
E0000 00:00:1743137370.408219 3481103 buffer_comparator.cc:156] Difference at 68: 10.1174, expected 9.00042
E0000 00:00:1743137370.408222 3481103 buffer_comparator.cc:156] Difference at 69: 9.89636, expected 8.345
E0000 00:00:1743137370.408225 3481103 buffer_comparator.cc:156] Difference at 70: 10.5296, expected 8.52479
E0000 00:00:1743137370.408228 3481103 buffer_comparator.cc:156] Difference at 71: 8.1412, expected 7.08541
E0000 00:00:1743137370.408231 3481103 buffer_comparator.cc:156] Difference at 72: 6.85179, expected 8.47292
E0000 00:00:1743137370.408234 3481103 buffer_comparator.cc:156] Difference at 77: 7.45304, expected 9.28323
E0000 00:00:1743137370.408236 3481103 buffer_comparator.cc:156] Difference at 80: 7.39802, expected 9.4894
E0000 00:00:1743137370.408239 3481103 buffer_comparator.cc:156] Difference at 81: 9.68936, expected 7.49939
E0000 00:00:1743137370.408242 3481103 buffer_comparator.cc:156] Difference at 82: 9.2677, expected 6.28629
2025-03-28 04:49:30.408246: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.410263 3481103 buffer_comparator.cc:156] Difference at 67: 10.5021, expected 7.70695
E0000 00:00:1743137370.410275 3481103 buffer_comparator.cc:156] Difference at 68: 10.1174, expected 9.00042
E0000 00:00:1743137370.410278 3481103 buffer_comparator.cc:156] Difference at 69: 9.89636, expected 8.345
E0000 00:00:1743137370.410281 3481103 buffer_comparator.cc:156] Difference at 70: 10.5296, expected 8.52479
E0000 00:00:1743137370.410284 3481103 buffer_comparator.cc:156] Difference at 71: 8.1412, expected 7.08541
E0000 00:00:1743137370.410287 3481103 buffer_comparator.cc:156] Difference at 72: 6.85179, expected 8.47292
E0000 00:00:1743137370.410290 3481103 buffer_comparator.cc:156] Difference at 77: 7.45304, expected 9.28323
E0000 00:00:1743137370.410293 3481103 buffer_comparator.cc:156] Difference at 80: 7.39802, expected 9.4894
E0000 00:00:1743137370.410296 3481103 buffer_comparator.cc:156] Difference at 81: 9.68936, expected 7.49939
E0000 00:00:1743137370.410298 3481103 buffer_comparator.cc:156] Difference at 82: 9.2677, expected 6.28629
2025-03-28 04:49:30.410303: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.411958 3481103 buffer_comparator.cc:156] Difference at 128: 8.2391, expected 10.3054
E0000 00:00:1743137370.411970 3481103 buffer_comparator.cc:156] Difference at 129: 8.47949, expected 9.87733
E0000 00:00:1743137370.411973 3481103 buffer_comparator.cc:156] Difference at 133: 8.18623, expected 9.79527
E0000 00:00:1743137370.411976 3481103 buffer_comparator.cc:156] Difference at 138: 6.64546, expected 8.48491
E0000 00:00:1743137370.411980 3481103 buffer_comparator.cc:156] Difference at 139: 8.57118, expected 7.35965
E0000 00:00:1743137370.411983 3481103 buffer_comparator.cc:156] Difference at 141: 8.57626, expected 7.30194
E0000 00:00:1743137370.411986 3481103 buffer_comparator.cc:156] Difference at 142: 8.50319, expected 7.45738
E0000 00:00:1743137370.411989 3481103 buffer_comparator.cc:156] Difference at 144: 8.42517, expected 7.38918
E0000 00:00:1743137370.411992 3481103 buffer_comparator.cc:156] Difference at 146: 9.43345, expected 7.98069
E0000 00:00:1743137370.411994 3481103 buffer_comparator.cc:156] Difference at 147: 7.98762, expected 9.22175
2025-03-28 04:49:30.411999: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.413747 3481103 buffer_comparator.cc:156] Difference at 128: 8.2391, expected 10.3054
E0000 00:00:1743137370.413759 3481103 buffer_comparator.cc:156] Difference at 129: 8.47949, expected 9.87733
E0000 00:00:1743137370.413762 3481103 buffer_comparator.cc:156] Difference at 133: 8.18623, expected 9.79527
E0000 00:00:1743137370.413765 3481103 buffer_comparator.cc:156] Difference at 138: 6.64546, expected 8.48491
E0000 00:00:1743137370.413768 3481103 buffer_comparator.cc:156] Difference at 139: 8.57118, expected 7.35965
E0000 00:00:1743137370.413770 3481103 buffer_comparator.cc:156] Difference at 141: 8.57626, expected 7.30194
E0000 00:00:1743137370.413773 3481103 buffer_comparator.cc:156] Difference at 142: 8.50319, expected 7.45738
E0000 00:00:1743137370.413776 3481103 buffer_comparator.cc:156] Difference at 144: 8.42517, expected 7.38918
E0000 00:00:1743137370.413779 3481103 buffer_comparator.cc:156] Difference at 146: 9.43345, expected 7.98069
E0000 00:00:1743137370.413782 3481103 buffer_comparator.cc:156] Difference at 147: 7.98762, expected 9.22175
2025-03-28 04:49:30.413786: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.419594 3481103 buffer_comparator.cc:156] Difference at 16: -nan, expected 8.81126
E0000 00:00:1743137370.419609 3481103 buffer_comparator.cc:156] Difference at 17: -nan, expected 10.499
E0000 00:00:1743137370.419613 3481103 buffer_comparator.cc:156] Difference at 18: -nan, expected 9.85809
E0000 00:00:1743137370.419615 3481103 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.65377
E0000 00:00:1743137370.419618 3481103 buffer_comparator.cc:156] Difference at 20: -nan, expected 10.0279
E0000 00:00:1743137370.419621 3481103 buffer_comparator.cc:156] Difference at 21: -nan, expected 9.616
E0000 00:00:1743137370.419624 3481103 buffer_comparator.cc:156] Difference at 22: -nan, expected 11.4222
E0000 00:00:1743137370.419627 3481103 buffer_comparator.cc:156] Difference at 23: -nan, expected 8.99654
E0000 00:00:1743137370.419629 3481103 buffer_comparator.cc:156] Difference at 24: -nan, expected 10.1518
E0000 00:00:1743137370.419632 3481103 buffer_comparator.cc:156] Difference at 25: -nan, expected 9.78295
2025-03-28 04:49:30.419637: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.421290 3481103 buffer_comparator.cc:156] Difference at 16: -nan, expected 8.81126
E0000 00:00:1743137370.421301 3481103 buffer_comparator.cc:156] Difference at 17: -nan, expected 10.499
E0000 00:00:1743137370.421305 3481103 buffer_comparator.cc:156] Difference at 18: -nan, expected 9.85809
E0000 00:00:1743137370.421307 3481103 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.65377
E0000 00:00:1743137370.421310 3481103 buffer_comparator.cc:156] Difference at 20: -nan, expected 10.0279
E0000 00:00:1743137370.421313 3481103 buffer_comparator.cc:156] Difference at 21: -nan, expected 9.616
E0000 00:00:1743137370.421316 3481103 buffer_comparator.cc:156] Difference at 22: -nan, expected 11.4222
E0000 00:00:1743137370.421320 3481103 buffer_comparator.cc:156] Difference at 23: -nan, expected 8.99654
E0000 00:00:1743137370.421323 3481103 buffer_comparator.cc:156] Difference at 24: -nan, expected 10.1518
E0000 00:00:1743137370.421326 3481103 buffer_comparator.cc:156] Difference at 25: -nan, expected 9.78295
2025-03-28 04:49:30.421330: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.422977 3481103 buffer_comparator.cc:156] Difference at 16: -nan, expected 8.81126
E0000 00:00:1743137370.422989 3481103 buffer_comparator.cc:156] Difference at 17: -nan, expected 10.499
E0000 00:00:1743137370.422992 3481103 buffer_comparator.cc:156] Difference at 18: -nan, expected 9.85809
E0000 00:00:1743137370.422995 3481103 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.65377
E0000 00:00:1743137370.422997 3481103 buffer_comparator.cc:156] Difference at 20: -nan, expected 10.0279
E0000 00:00:1743137370.423000 3481103 buffer_comparator.cc:156] Difference at 21: -nan, expected 9.616
E0000 00:00:1743137370.423003 3481103 buffer_comparator.cc:156] Difference at 22: -nan, expected 11.4222
E0000 00:00:1743137370.423006 3481103 buffer_comparator.cc:156] Difference at 23: -nan, expected 8.99654
E0000 00:00:1743137370.423008 3481103 buffer_comparator.cc:156] Difference at 24: -nan, expected 10.1518
E0000 00:00:1743137370.423011 3481103 buffer_comparator.cc:156] Difference at 25: -nan, expected 9.78295
2025-03-28 04:49:30.423015: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.424670 3481103 buffer_comparator.cc:156] Difference at 32: -nan, expected 8.20589
E0000 00:00:1743137370.424681 3481103 buffer_comparator.cc:156] Difference at 33: -nan, expected 9.19448
E0000 00:00:1743137370.424684 3481103 buffer_comparator.cc:156] Difference at 34: -nan, expected 8.03576
E0000 00:00:1743137370.424687 3481103 buffer_comparator.cc:156] Difference at 35: -nan, expected 10.6612
E0000 00:00:1743137370.424690 3481103 buffer_comparator.cc:156] Difference at 36: -nan, expected 9.50115
E0000 00:00:1743137370.424692 3481103 buffer_comparator.cc:156] Difference at 37: -nan, expected 10.8965
E0000 00:00:1743137370.424695 3481103 buffer_comparator.cc:156] Difference at 38: -nan, expected 9.00596
E0000 00:00:1743137370.424698 3481103 buffer_comparator.cc:156] Difference at 39: -nan, expected 7.25269
E0000 00:00:1743137370.424701 3481103 buffer_comparator.cc:156] Difference at 40: -nan, expected 9.67287
E0000 00:00:1743137370.424703 3481103 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.30318
2025-03-28 04:49:30.424708: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.426350 3481103 buffer_comparator.cc:156] Difference at 32: -nan, expected 8.20589
E0000 00:00:1743137370.426361 3481103 buffer_comparator.cc:156] Difference at 33: -nan, expected 9.19448
E0000 00:00:1743137370.426364 3481103 buffer_comparator.cc:156] Difference at 34: -nan, expected 8.03576
E0000 00:00:1743137370.426367 3481103 buffer_comparator.cc:156] Difference at 35: -nan, expected 10.6612
E0000 00:00:1743137370.426370 3481103 buffer_comparator.cc:156] Difference at 36: -nan, expected 9.50115
E0000 00:00:1743137370.426372 3481103 buffer_comparator.cc:156] Difference at 37: -nan, expected 10.8965
E0000 00:00:1743137370.426375 3481103 buffer_comparator.cc:156] Difference at 38: -nan, expected 9.00596
E0000 00:00:1743137370.426378 3481103 buffer_comparator.cc:156] Difference at 39: -nan, expected 7.25269
E0000 00:00:1743137370.426381 3481103 buffer_comparator.cc:156] Difference at 40: -nan, expected 9.67287
E0000 00:00:1743137370.426383 3481103 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.30318
2025-03-28 04:49:30.426390: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1743137370.428038 3481103 buffer_comparator.cc:156] Difference at 32: -nan, expected 8.20589
E0000 00:00:1743137370.428049 3481103 buffer_comparator.cc:156] Difference at 33: -nan, expected 9.19448
E0000 00:00:1743137370.428052 3481103 buffer_comparator.cc:156] Difference at 34: -nan, expected 8.03576
E0000 00:00:1743137370.428055 3481103 buffer_comparator.cc:156] Difference at 35: -nan, expected 10.6612
E0000 00:00:1743137370.428058 3481103 buffer_comparator.cc:156] Difference at 36: -nan, expected 9.50115
E0000 00:00:1743137370.428061 3481103 buffer_comparator.cc:156] Difference at 37: -nan, expected 10.8965
E0000 00:00:1743137370.428064 3481103 buffer_comparator.cc:156] Difference at 38: -nan, expected 9.00596
E0000 00:00:1743137370.428066 3481103 buffer_comparator.cc:156] Difference at 39: -nan, expected 7.25269
E0000 00:00:1743137370.428069 3481103 buffer_comparator.cc:156] Difference at 40: -nan, expected 9.67287
E0000 00:00:1743137370.428072 3481103 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.30318
2025-03-28 04:49:30.428076: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1137] Results do not match the reference. This is likely a bug/unexpected loss of precision.
Iteration: [     1/ 50000] 	 Loss: 3.153872967 (3.153872967) 	 Physics Loss: 1.977264166 (1.977264166) 	 Data Loss: 0.578239620 (0.578239620) 	 BC Loss: 0.598369122 (0.598369122)
Iteration: [  1001/ 50000] 	 Loss: 0.037203494 (0.029109152) 	 Physics Loss: 0.000550530 (0.000358969) 	 Data Loss: 0.023778778 (0.012030440) 	 BC Loss: 0.012874187 (0.016719742)
Iteration: [  2001/ 50000] 	 Loss: 0.030779511 (0.034423992) 	 Physics Loss: 0.000592207 (0.001774628) 	 Data Loss: 0.014397524 (0.013753930) 	 BC Loss: 0.015789781 (0.018895429)
Iteration: [  3001/ 50000] 	 Loss: 0.021841470 (0.027056616) 	 Physics Loss: 0.000522536 (0.004873221) 	 Data Loss: 0.004614574 (0.008088700) 	 BC Loss: 0.016704362 (0.014094692)
Iteration: [  4001/ 50000] 	 Loss: 0.038582101 (0.034729011) 	 Physics Loss: 0.003376254 (0.002704982) 	 Data Loss: 0.014301238 (0.012586679) 	 BC Loss: 0.020904612 (0.019437356)
Iteration: [  5001/ 50000] 	 Loss: 0.027363494 (0.057518166) 	 Physics Loss: 0.006317818 (0.028009484) 	 Data Loss: 0.009285189 (0.011741836) 	 BC Loss: 0.011760486 (0.017766841)
Iteration: [  6001/ 50000] 	 Loss: 0.024195662 (0.019290088) 	 Physics Loss: 0.000576330 (0.000937816) 	 Data Loss: 0.005862516 (0.006385402) 	 BC Loss: 0.017756816 (0.011966869)
Iteration: [  7001/ 50000] 	 Loss: 0.015304947 (0.019078610) 	 Physics Loss: 0.001520974 (0.001128703) 	 Data Loss: 0.005780942 (0.006937551) 	 BC Loss: 0.008003032 (0.011012357)
Iteration: [  8001/ 50000] 	 Loss: 0.018056616 (0.016851155) 	 Physics Loss: 0.000815929 (0.001745509) 	 Data Loss: 0.003489140 (0.004778590) 	 BC Loss: 0.013751546 (0.010327056)
Iteration: [  9001/ 50000] 	 Loss: 0.019105814 (0.016347572) 	 Physics Loss: 0.001103567 (0.001558701) 	 Data Loss: 0.002615690 (0.004747680) 	 BC Loss: 0.015386556 (0.010041193)
Iteration: [ 10001/ 50000] 	 Loss: 0.011850247 (0.016185638) 	 Physics Loss: 0.001527360 (0.002253429) 	 Data Loss: 0.004485745 (0.004699244) 	 BC Loss: 0.005837142 (0.009232965)
Iteration: [ 11001/ 50000] 	 Loss: 0.017559819 (0.014584152) 	 Physics Loss: 0.001138550 (0.001371698) 	 Data Loss: 0.004524786 (0.004185131) 	 BC Loss: 0.011896483 (0.009027325)
Iteration: [ 12001/ 50000] 	 Loss: 0.014435371 (0.013971340) 	 Physics Loss: 0.001502353 (0.001294284) 	 Data Loss: 0.006330506 (0.004238269) 	 BC Loss: 0.006602511 (0.008438786)
Iteration: [ 13001/ 50000] 	 Loss: 0.007569676 (0.013424525) 	 Physics Loss: 0.001393272 (0.001443758) 	 Data Loss: 0.001905998 (0.003799081) 	 BC Loss: 0.004270406 (0.008181685)
Iteration: [ 14001/ 50000] 	 Loss: 0.008399807 (0.012269501) 	 Physics Loss: 0.000874100 (0.001270473) 	 Data Loss: 0.003854776 (0.003514751) 	 BC Loss: 0.003670931 (0.007484275)
Iteration: [ 15001/ 50000] 	 Loss: 0.012854995 (0.012447622) 	 Physics Loss: 0.000962182 (0.001374702) 	 Data Loss: 0.000483413 (0.003976720) 	 BC Loss: 0.011409399 (0.007096201)
Iteration: [ 16001/ 50000] 	 Loss: 0.008248669 (0.013208234) 	 Physics Loss: 0.001395914 (0.001164462) 	 Data Loss: 0.003623560 (0.003468768) 	 BC Loss: 0.003229195 (0.008575005)
Iteration: [ 17001/ 50000] 	 Loss: 0.005129706 (0.012013515) 	 Physics Loss: 0.000699974 (0.001131699) 	 Data Loss: 0.000917939 (0.003342754) 	 BC Loss: 0.003511793 (0.007539063)
Iteration: [ 18001/ 50000] 	 Loss: 0.016537888 (0.010515220) 	 Physics Loss: 0.001252979 (0.001072636) 	 Data Loss: 0.004219135 (0.002489778) 	 BC Loss: 0.011065774 (0.006952806)
Iteration: [ 19001/ 50000] 	 Loss: 0.010813366 (0.012705814) 	 Physics Loss: 0.001188082 (0.001193895) 	 Data Loss: 0.003777510 (0.002822159) 	 BC Loss: 0.005847774 (0.008689760)
Iteration: [ 20001/ 50000] 	 Loss: 0.005427668 (0.011710118) 	 Physics Loss: 0.001577919 (0.001387226) 	 Data Loss: 0.002568154 (0.002168224) 	 BC Loss: 0.001281594 (0.008154668)
Iteration: [ 21001/ 50000] 	 Loss: 0.006487023 (0.010208706) 	 Physics Loss: 0.001401789 (0.001660767) 	 Data Loss: 0.001722868 (0.002501602) 	 BC Loss: 0.003362367 (0.006046339)
Iteration: [ 22001/ 50000] 	 Loss: 0.009437616 (0.010694695) 	 Physics Loss: 0.001765511 (0.001500460) 	 Data Loss: 0.005958611 (0.001754927) 	 BC Loss: 0.001713495 (0.007439307)
Iteration: [ 23001/ 50000] 	 Loss: 0.013743464 (0.008254445) 	 Physics Loss: 0.002180674 (0.001913983) 	 Data Loss: 0.004219715 (0.002065712) 	 BC Loss: 0.007343076 (0.004274749)
Iteration: [ 24001/ 50000] 	 Loss: 0.014256330 (0.008225464) 	 Physics Loss: 0.001869628 (0.001696485) 	 Data Loss: 0.000880123 (0.001554014) 	 BC Loss: 0.011506579 (0.004974965)
Iteration: [ 25001/ 50000] 	 Loss: 0.005411615 (0.006930466) 	 Physics Loss: 0.001179156 (0.001798108) 	 Data Loss: 0.000670032 (0.001461951) 	 BC Loss: 0.003562426 (0.003670407)
Iteration: [ 26001/ 50000] 	 Loss: 0.006880802 (0.005627883) 	 Physics Loss: 0.001134805 (0.001795200) 	 Data Loss: 0.000952239 (0.001306181) 	 BC Loss: 0.004793759 (0.002526500)
Iteration: [ 27001/ 50000] 	 Loss: 0.007057397 (0.005158458) 	 Physics Loss: 0.000736209 (0.001526923) 	 Data Loss: 0.001355446 (0.001198944) 	 BC Loss: 0.004965742 (0.002432592)
Iteration: [ 28001/ 50000] 	 Loss: 0.002918471 (0.004430990) 	 Physics Loss: 0.001545964 (0.001843499) 	 Data Loss: 0.000570738 (0.001045180) 	 BC Loss: 0.000801769 (0.001542311)
Iteration: [ 29001/ 50000] 	 Loss: 0.003383059 (0.003222783) 	 Physics Loss: 0.001546039 (0.001197281) 	 Data Loss: 0.001057410 (0.000871451) 	 BC Loss: 0.000779610 (0.001154052)
Iteration: [ 30001/ 50000] 	 Loss: 0.003592887 (0.003187301) 	 Physics Loss: 0.001639638 (0.001459631) 	 Data Loss: 0.000699877 (0.000813704) 	 BC Loss: 0.001253372 (0.000913966)
Iteration: [ 31001/ 50000] 	 Loss: 0.002682653 (0.002753643) 	 Physics Loss: 0.000679309 (0.001226370) 	 Data Loss: 0.000449619 (0.000731408) 	 BC Loss: 0.001553725 (0.000795864)
Iteration: [ 32001/ 50000] 	 Loss: 0.005353532 (0.003093726) 	 Physics Loss: 0.003998624 (0.001627845) 	 Data Loss: 0.000485983 (0.000814618) 	 BC Loss: 0.000868925 (0.000651262)
Iteration: [ 33001/ 50000] 	 Loss: 0.001863127 (0.002165305) 	 Physics Loss: 0.000551163 (0.001070478) 	 Data Loss: 0.000602129 (0.000578883) 	 BC Loss: 0.000709834 (0.000515945)
Iteration: [ 34001/ 50000] 	 Loss: 0.001789484 (0.002187706) 	 Physics Loss: 0.000852895 (0.001082061) 	 Data Loss: 0.000387473 (0.000637523) 	 BC Loss: 0.000549116 (0.000468122)
Iteration: [ 35001/ 50000] 	 Loss: 0.000949289 (0.001767912) 	 Physics Loss: 0.000569952 (0.000929367) 	 Data Loss: 0.000198365 (0.000504372) 	 BC Loss: 0.000180972 (0.000334173)
Iteration: [ 36001/ 50000] 	 Loss: 0.002945340 (0.001718901) 	 Physics Loss: 0.001339663 (0.000954814) 	 Data Loss: 0.001226619 (0.000454665) 	 BC Loss: 0.000379057 (0.000309422)
Iteration: [ 37001/ 50000] 	 Loss: 0.001214534 (0.001479726) 	 Physics Loss: 0.000422838 (0.000726230) 	 Data Loss: 0.000534357 (0.000504710) 	 BC Loss: 0.000257339 (0.000248787)
Iteration: [ 38001/ 50000] 	 Loss: 0.001880942 (0.001397801) 	 Physics Loss: 0.000924640 (0.000684149) 	 Data Loss: 0.000530066 (0.000463928) 	 BC Loss: 0.000426236 (0.000249725)
Iteration: [ 39001/ 50000] 	 Loss: 0.001347756 (0.001421823) 	 Physics Loss: 0.000728663 (0.000679449) 	 Data Loss: 0.000345473 (0.000504370) 	 BC Loss: 0.000273620 (0.000238003)
Iteration: [ 40001/ 50000] 	 Loss: 0.001802240 (0.001479169) 	 Physics Loss: 0.001196385 (0.000767612) 	 Data Loss: 0.000447591 (0.000483108) 	 BC Loss: 0.000158264 (0.000228449)
Iteration: [ 41001/ 50000] 	 Loss: 0.001352749 (0.001697045) 	 Physics Loss: 0.000932642 (0.000977645) 	 Data Loss: 0.000314114 (0.000515789) 	 BC Loss: 0.000105992 (0.000203611)
Iteration: [ 42001/ 50000] 	 Loss: 0.001608718 (0.001346420) 	 Physics Loss: 0.001242280 (0.000718649) 	 Data Loss: 0.000229054 (0.000422413) 	 BC Loss: 0.000137383 (0.000205358)
Iteration: [ 43001/ 50000] 	 Loss: 0.000952899 (0.001501483) 	 Physics Loss: 0.000515131 (0.000807861) 	 Data Loss: 0.000285539 (0.000429003) 	 BC Loss: 0.000152229 (0.000264619)
Iteration: [ 44001/ 50000] 	 Loss: 0.000931397 (0.001297562) 	 Physics Loss: 0.000534546 (0.000711467) 	 Data Loss: 0.000204832 (0.000406957) 	 BC Loss: 0.000192019 (0.000179138)
Iteration: [ 45001/ 50000] 	 Loss: 0.001132066 (0.001029962) 	 Physics Loss: 0.000462024 (0.000488499) 	 Data Loss: 0.000538058 (0.000379111) 	 BC Loss: 0.000131983 (0.000162352)
Iteration: [ 46001/ 50000] 	 Loss: 0.001625891 (0.001669809) 	 Physics Loss: 0.001002158 (0.001080601) 	 Data Loss: 0.000345248 (0.000392433) 	 BC Loss: 0.000278485 (0.000196775)
Iteration: [ 47001/ 50000] 	 Loss: 0.001482891 (0.001004653) 	 Physics Loss: 0.000928177 (0.000492962) 	 Data Loss: 0.000446164 (0.000375494) 	 BC Loss: 0.000108550 (0.000136197)
Iteration: [ 48001/ 50000] 	 Loss: 0.001417234 (0.001158843) 	 Physics Loss: 0.000884422 (0.000662071) 	 Data Loss: 0.000195326 (0.000363015) 	 BC Loss: 0.000337487 (0.000133758)
Iteration: [ 49001/ 50000] 	 Loss: 0.001132041 (0.001062525) 	 Physics Loss: 0.000717354 (0.000549186) 	 Data Loss: 0.000253737 (0.000361861) 	 BC Loss: 0.000160949 (0.000151479)

Visualizing the Results

julia

ts, xs, ys = 0.0f0:0.05f0:2.0f0, 0.0f0:0.02f0:2.0f0, 0.0f0:0.02f0:2.0f0
grid = stack([[elem...] for elem in vec(collect(Iterators.product(xs, ys, ts)))])

u_real = reshape(analytical_solution(grid), length(xs), length(ys), length(ts))

grid_normalized = (grid .- minimum(grid)) ./ (maximum(grid) .- minimum(grid))
u_pred = reshape(trained_u(grid_normalized), length(xs), length(ys), length(ts))
u_pred = u_pred .* (max_pde_val - min_pde_val) .+ min_pde_val

begin
    fig = Figure()
    ax = CairoMakie.Axis(fig[1, 1]; xlabel="x", ylabel="y")
    errs = [abs.(u_pred[:, :, i] .- u_real[:, :, i]) for i in 1:length(ts)]
    Colorbar(fig[1, 2]; limits=extrema(stack(errs)))

    CairoMakie.record(fig, "pinn_nested_ad.gif", 1:length(ts); framerate=10) do i
        ax.title = "Abs. Predictor Error | Time: $(ts[i])"
        err = errs[i]
        contour!(ax, xs, ys, err; levels=10, linewidth=2)
        heatmap!(ax, xs, ys, err)
        return fig
    end

    fig
end

Appendix

julia

using InteractiveUtils
InteractiveUtils.versioninfo()

if @isdefined(MLDataDevices)
    if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice)
        println()
        CUDA.versioninfo()
    end

    if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice)
        println()
        AMDGPU.versioninfo()
    end
end

Julia Version 1.11.4
Commit 8561cc3d68d (2025-03-10 11:36 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 48 × AMD EPYC 7402 24-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver2)
Threads: 48 default, 0 interactive, 24 GC (on 2 virtual cores)
Environment:
  JULIA_CPU_THREADS = 2
  LD_LIBRARY_PATH = /usr/local/nvidia/lib:/usr/local/nvidia/lib64
  JULIA_PKG_SERVER = 
  JULIA_NUM_THREADS = 48
  JULIA_CUDA_HARD_MEMORY_LIMIT = 100%
  JULIA_PKG_PRECOMPILE_AUTO = 0
  JULIA_DEBUG = Literate
  JULIA_DEPOT_PATH = /root/.cache/julia-buildkite-plugin/depots/01872db4-8c79-43af-ab7d-12abac4f24f6

This page was generated using Literate.jl.

Trusted by

Training a PINN on 2D PDE ​

Package Imports ​

Problem Definition ​

Define the Neural Networks ​

Define the Loss Functions ​

Generate the Data ​

Training ​

Visualizing the Results ​

Appendix ​

Training a PINN on 2D PDE

Package Imports

Problem Definition

Define the Neural Networks

Define the Loss Functions

Generate the Data

Training

Visualizing the Results

Appendix