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 = Lux.setup(rng, pinn) |> xdev

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

    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;
            return_gradients=Val(false),
        )

        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-04-07 02:09:49.679298: I external/xla/xla/service/service.cc:152] XLA service 0xb70a840 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2025-04-07 02:09:49.679338: 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:1743991789.680104 3863697 se_gpu_pjrt_client.cc:1040] Using BFC allocator.
I0000 00:00:1743991789.680166 3863697 gpu_helpers.cc:136] XLA backend allocating 3825205248 bytes on device 0 for BFCAllocator.
I0000 00:00:1743991789.680209 3863697 gpu_helpers.cc:177] XLA backend will use up to 1275068416 bytes on device 0 for CollectiveBFCAllocator.
I0000 00:00:1743991789.692540 3863697 cuda_dnn.cc:529] Loaded cuDNN version 90400
E0000 00:00:1743992096.874632 3863697 buffer_comparator.cc:156] Difference at 16: 0, expected 9.278
E0000 00:00:1743992096.874680 3863697 buffer_comparator.cc:156] Difference at 17: 0, expected 9.69724
E0000 00:00:1743992096.874688 3863697 buffer_comparator.cc:156] Difference at 18: 0, expected 10.7774
E0000 00:00:1743992096.874696 3863697 buffer_comparator.cc:156] Difference at 19: 0, expected 10.3636
E0000 00:00:1743992096.874703 3863697 buffer_comparator.cc:156] Difference at 20: 0, expected 7.56533
E0000 00:00:1743992096.874709 3863697 buffer_comparator.cc:156] Difference at 21: 0, expected 8.17756
E0000 00:00:1743992096.874715 3863697 buffer_comparator.cc:156] Difference at 22: 0, expected 11.0551
E0000 00:00:1743992096.874722 3863697 buffer_comparator.cc:156] Difference at 23: 0, expected 11.628
E0000 00:00:1743992096.874729 3863697 buffer_comparator.cc:156] Difference at 24: 0, expected 10.6715
E0000 00:00:1743992096.874735 3863697 buffer_comparator.cc:156] Difference at 25: 0, expected 10.8131
2025-04-07 02:14:56.874750: 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:1743992096.877486 3863697 buffer_comparator.cc:156] Difference at 16: 0, expected 9.278
E0000 00:00:1743992096.877510 3863697 buffer_comparator.cc:156] Difference at 17: 0, expected 9.69724
E0000 00:00:1743992096.877517 3863697 buffer_comparator.cc:156] Difference at 18: 0, expected 10.7774
E0000 00:00:1743992096.877524 3863697 buffer_comparator.cc:156] Difference at 19: 0, expected 10.3636
E0000 00:00:1743992096.877531 3863697 buffer_comparator.cc:156] Difference at 20: 0, expected 7.56533
E0000 00:00:1743992096.877537 3863697 buffer_comparator.cc:156] Difference at 21: 0, expected 8.17756
E0000 00:00:1743992096.877543 3863697 buffer_comparator.cc:156] Difference at 22: 0, expected 11.0551
E0000 00:00:1743992096.877550 3863697 buffer_comparator.cc:156] Difference at 23: 0, expected 11.628
E0000 00:00:1743992096.877556 3863697 buffer_comparator.cc:156] Difference at 24: 0, expected 10.6715
E0000 00:00:1743992096.877563 3863697 buffer_comparator.cc:156] Difference at 25: 0, expected 10.8131
2025-04-07 02:14:56.877574: 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:1743992096.880026 3863697 buffer_comparator.cc:156] Difference at 16: 0, expected 9.278
E0000 00:00:1743992096.880050 3863697 buffer_comparator.cc:156] Difference at 17: 0, expected 9.69724
E0000 00:00:1743992096.880057 3863697 buffer_comparator.cc:156] Difference at 18: 0, expected 10.7774
E0000 00:00:1743992096.880064 3863697 buffer_comparator.cc:156] Difference at 19: 0, expected 10.3636
E0000 00:00:1743992096.880070 3863697 buffer_comparator.cc:156] Difference at 20: 0, expected 7.56533
E0000 00:00:1743992096.880077 3863697 buffer_comparator.cc:156] Difference at 21: 0, expected 8.17756
E0000 00:00:1743992096.880083 3863697 buffer_comparator.cc:156] Difference at 22: 0, expected 11.0551
E0000 00:00:1743992096.880089 3863697 buffer_comparator.cc:156] Difference at 23: 0, expected 11.628
E0000 00:00:1743992096.880096 3863697 buffer_comparator.cc:156] Difference at 24: 0, expected 10.6715
E0000 00:00:1743992096.880105 3863697 buffer_comparator.cc:156] Difference at 25: 0, expected 10.8131
2025-04-07 02:14:56.880116: 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:1743992096.882554 3863697 buffer_comparator.cc:156] Difference at 32: 0, expected 10.5897
E0000 00:00:1743992096.882579 3863697 buffer_comparator.cc:156] Difference at 33: 0, expected 8.79857
E0000 00:00:1743992096.882586 3863697 buffer_comparator.cc:156] Difference at 34: 0, expected 10.7009
E0000 00:00:1743992096.882593 3863697 buffer_comparator.cc:156] Difference at 35: 0, expected 10.0423
E0000 00:00:1743992096.882599 3863697 buffer_comparator.cc:156] Difference at 36: 0, expected 10.3055
E0000 00:00:1743992096.882605 3863697 buffer_comparator.cc:156] Difference at 37: 0, expected 11.1027
E0000 00:00:1743992096.882612 3863697 buffer_comparator.cc:156] Difference at 38: 0, expected 9.64939
E0000 00:00:1743992096.882618 3863697 buffer_comparator.cc:156] Difference at 39: 0, expected 10.2558
E0000 00:00:1743992096.882624 3863697 buffer_comparator.cc:156] Difference at 40: 0, expected 8.38055
E0000 00:00:1743992096.882631 3863697 buffer_comparator.cc:156] Difference at 41: 0, expected 7.92221
2025-04-07 02:14:56.882641: 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:1743992096.885064 3863697 buffer_comparator.cc:156] Difference at 32: 0, expected 10.5897
E0000 00:00:1743992096.885089 3863697 buffer_comparator.cc:156] Difference at 33: 0, expected 8.79857
E0000 00:00:1743992096.885096 3863697 buffer_comparator.cc:156] Difference at 34: 0, expected 10.7009
E0000 00:00:1743992096.885102 3863697 buffer_comparator.cc:156] Difference at 35: 0, expected 10.0423
E0000 00:00:1743992096.885109 3863697 buffer_comparator.cc:156] Difference at 36: 0, expected 10.3055
E0000 00:00:1743992096.885115 3863697 buffer_comparator.cc:156] Difference at 37: 0, expected 11.1027
E0000 00:00:1743992096.885122 3863697 buffer_comparator.cc:156] Difference at 38: 0, expected 9.64939
E0000 00:00:1743992096.885128 3863697 buffer_comparator.cc:156] Difference at 39: 0, expected 10.2558
E0000 00:00:1743992096.885134 3863697 buffer_comparator.cc:156] Difference at 40: 0, expected 8.38055
E0000 00:00:1743992096.885140 3863697 buffer_comparator.cc:156] Difference at 41: 0, expected 7.92221
2025-04-07 02:14:56.885150: 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:1743992096.887537 3863697 buffer_comparator.cc:156] Difference at 32: 0, expected 10.5897
E0000 00:00:1743992096.887550 3863697 buffer_comparator.cc:156] Difference at 33: 0, expected 8.79857
E0000 00:00:1743992096.887553 3863697 buffer_comparator.cc:156] Difference at 34: 0, expected 10.7009
E0000 00:00:1743992096.887556 3863697 buffer_comparator.cc:156] Difference at 35: 0, expected 10.0423
E0000 00:00:1743992096.887559 3863697 buffer_comparator.cc:156] Difference at 36: 0, expected 10.3055
E0000 00:00:1743992096.887561 3863697 buffer_comparator.cc:156] Difference at 37: 0, expected 11.1027
E0000 00:00:1743992096.887564 3863697 buffer_comparator.cc:156] Difference at 38: 0, expected 9.64939
E0000 00:00:1743992096.887567 3863697 buffer_comparator.cc:156] Difference at 39: 0, expected 10.2558
E0000 00:00:1743992096.887570 3863697 buffer_comparator.cc:156] Difference at 40: 0, expected 8.38055
E0000 00:00:1743992096.887573 3863697 buffer_comparator.cc:156] Difference at 41: 0, expected 7.92221
2025-04-07 02:14:56.887577: 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:1743992096.889700 3863697 buffer_comparator.cc:156] Difference at 64: 0, expected 8.4725
E0000 00:00:1743992096.889713 3863697 buffer_comparator.cc:156] Difference at 65: 0, expected 8.00729
E0000 00:00:1743992096.889716 3863697 buffer_comparator.cc:156] Difference at 66: 0, expected 10.5558
E0000 00:00:1743992096.889719 3863697 buffer_comparator.cc:156] Difference at 67: 0, expected 7.97013
E0000 00:00:1743992096.889722 3863697 buffer_comparator.cc:156] Difference at 68: 0, expected 9.17677
E0000 00:00:1743992096.889724 3863697 buffer_comparator.cc:156] Difference at 69: 0, expected 6.91094
E0000 00:00:1743992096.889727 3863697 buffer_comparator.cc:156] Difference at 70: 0, expected 7.85311
E0000 00:00:1743992096.889730 3863697 buffer_comparator.cc:156] Difference at 71: 0, expected 8.69839
E0000 00:00:1743992096.889733 3863697 buffer_comparator.cc:156] Difference at 72: 0, expected 9.37233
E0000 00:00:1743992096.889735 3863697 buffer_comparator.cc:156] Difference at 73: 0, expected 7.56939
2025-04-07 02:14:56.889740: 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:1743992096.891879 3863697 buffer_comparator.cc:156] Difference at 64: 0, expected 8.4725
E0000 00:00:1743992096.891890 3863697 buffer_comparator.cc:156] Difference at 65: 0, expected 8.00729
E0000 00:00:1743992096.891893 3863697 buffer_comparator.cc:156] Difference at 66: 0, expected 10.5558
E0000 00:00:1743992096.891896 3863697 buffer_comparator.cc:156] Difference at 67: 0, expected 7.97013
E0000 00:00:1743992096.891899 3863697 buffer_comparator.cc:156] Difference at 68: 0, expected 9.17677
E0000 00:00:1743992096.891901 3863697 buffer_comparator.cc:156] Difference at 69: 0, expected 6.91094
E0000 00:00:1743992096.891904 3863697 buffer_comparator.cc:156] Difference at 70: 0, expected 7.85311
E0000 00:00:1743992096.891907 3863697 buffer_comparator.cc:156] Difference at 71: 0, expected 8.69839
E0000 00:00:1743992096.891910 3863697 buffer_comparator.cc:156] Difference at 72: 0, expected 9.37233
E0000 00:00:1743992096.891912 3863697 buffer_comparator.cc:156] Difference at 73: 0, expected 7.56939
2025-04-07 02:14:56.891917: 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:1743992096.894042 3863697 buffer_comparator.cc:156] Difference at 64: 0, expected 8.4725
E0000 00:00:1743992096.894058 3863697 buffer_comparator.cc:156] Difference at 65: 0, expected 8.00729
E0000 00:00:1743992096.894061 3863697 buffer_comparator.cc:156] Difference at 66: 0, expected 10.5558
E0000 00:00:1743992096.894064 3863697 buffer_comparator.cc:156] Difference at 67: 0, expected 7.97013
E0000 00:00:1743992096.894066 3863697 buffer_comparator.cc:156] Difference at 68: 0, expected 9.17677
E0000 00:00:1743992096.894069 3863697 buffer_comparator.cc:156] Difference at 69: 0, expected 6.91094
E0000 00:00:1743992096.894072 3863697 buffer_comparator.cc:156] Difference at 70: 0, expected 7.85311
E0000 00:00:1743992096.894075 3863697 buffer_comparator.cc:156] Difference at 71: 0, expected 8.69839
E0000 00:00:1743992096.894077 3863697 buffer_comparator.cc:156] Difference at 72: 0, expected 9.37233
E0000 00:00:1743992096.894080 3863697 buffer_comparator.cc:156] Difference at 73: 0, expected 7.56939
2025-04-07 02:14:56.894085: 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:1743992096.901639 3863697 buffer_comparator.cc:156] Difference at 16: 7.75946, expected 8.81126
E0000 00:00:1743992096.901650 3863697 buffer_comparator.cc:156] Difference at 17: 8.06305, expected 10.499
E0000 00:00:1743992096.901654 3863697 buffer_comparator.cc:156] Difference at 18: 7.28087, expected 9.85809
E0000 00:00:1743992096.901657 3863697 buffer_comparator.cc:156] Difference at 20: 8.32296, expected 10.0279
E0000 00:00:1743992096.901660 3863697 buffer_comparator.cc:156] Difference at 21: 6.36154, expected 9.616
E0000 00:00:1743992096.901665 3863697 buffer_comparator.cc:156] Difference at 22: 6.2616, expected 11.4222
E0000 00:00:1743992096.901668 3863697 buffer_comparator.cc:156] Difference at 23: 7.0749, expected 8.99654
E0000 00:00:1743992096.901671 3863697 buffer_comparator.cc:156] Difference at 24: 6.97031, expected 10.1518
E0000 00:00:1743992096.901674 3863697 buffer_comparator.cc:156] Difference at 25: 7.34104, expected 9.78295
E0000 00:00:1743992096.901676 3863697 buffer_comparator.cc:156] Difference at 26: 6.30294, expected 11.1019
2025-04-07 02:14:56.901681: 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:1743992096.903813 3863697 buffer_comparator.cc:156] Difference at 16: 7.75946, expected 8.81126
E0000 00:00:1743992096.903824 3863697 buffer_comparator.cc:156] Difference at 17: 8.06305, expected 10.499
E0000 00:00:1743992096.903827 3863697 buffer_comparator.cc:156] Difference at 18: 7.28087, expected 9.85809
E0000 00:00:1743992096.903830 3863697 buffer_comparator.cc:156] Difference at 20: 8.32296, expected 10.0279
E0000 00:00:1743992096.903833 3863697 buffer_comparator.cc:156] Difference at 21: 6.36154, expected 9.616
E0000 00:00:1743992096.903836 3863697 buffer_comparator.cc:156] Difference at 22: 6.2616, expected 11.4222
E0000 00:00:1743992096.903839 3863697 buffer_comparator.cc:156] Difference at 23: 7.0749, expected 8.99654
E0000 00:00:1743992096.903842 3863697 buffer_comparator.cc:156] Difference at 24: 6.97031, expected 10.1518
E0000 00:00:1743992096.903844 3863697 buffer_comparator.cc:156] Difference at 25: 7.34104, expected 9.78295
E0000 00:00:1743992096.903847 3863697 buffer_comparator.cc:156] Difference at 26: 6.30294, expected 11.1019
2025-04-07 02:14:56.903852: 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:1743992096.905971 3863697 buffer_comparator.cc:156] Difference at 16: 7.75946, expected 8.81126
E0000 00:00:1743992096.905981 3863697 buffer_comparator.cc:156] Difference at 17: 8.06305, expected 10.499
E0000 00:00:1743992096.905985 3863697 buffer_comparator.cc:156] Difference at 18: 7.28087, expected 9.85809
E0000 00:00:1743992096.905987 3863697 buffer_comparator.cc:156] Difference at 20: 8.32296, expected 10.0279
E0000 00:00:1743992096.905990 3863697 buffer_comparator.cc:156] Difference at 21: 6.36154, expected 9.616
E0000 00:00:1743992096.905993 3863697 buffer_comparator.cc:156] Difference at 22: 6.2616, expected 11.4222
E0000 00:00:1743992096.905996 3863697 buffer_comparator.cc:156] Difference at 23: 7.0749, expected 8.99654
E0000 00:00:1743992096.905999 3863697 buffer_comparator.cc:156] Difference at 24: 6.97031, expected 10.1518
E0000 00:00:1743992096.906002 3863697 buffer_comparator.cc:156] Difference at 25: 7.34104, expected 9.78295
E0000 00:00:1743992096.906004 3863697 buffer_comparator.cc:156] Difference at 26: 6.30294, expected 11.1019
2025-04-07 02:14:56.906009: 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:1743992096.908139 3863697 buffer_comparator.cc:156] Difference at 35: 8.74559, expected 10.6612
E0000 00:00:1743992096.908150 3863697 buffer_comparator.cc:156] Difference at 37: 8.80435, expected 10.8965
E0000 00:00:1743992096.908153 3863697 buffer_comparator.cc:156] Difference at 41: 8.5867, expected 7.30318
E0000 00:00:1743992096.908156 3863697 buffer_comparator.cc:156] Difference at 42: 6.60885, expected 9.79879
E0000 00:00:1743992096.908159 3863697 buffer_comparator.cc:156] Difference at 43: 8.36889, expected 9.55385
E0000 00:00:1743992096.908162 3863697 buffer_comparator.cc:156] Difference at 46: 7.80689, expected 9.75357
E0000 00:00:1743992096.908165 3863697 buffer_comparator.cc:156] Difference at 47: 8.61022, expected 9.77022
E0000 00:00:1743992096.908170 3863697 buffer_comparator.cc:156] Difference at 48: 8.38568, expected 9.70077
E0000 00:00:1743992096.908173 3863697 buffer_comparator.cc:156] Difference at 50: 9.92471, expected 11.4143
E0000 00:00:1743992096.908176 3863697 buffer_comparator.cc:156] Difference at 51: 7.97977, expected 10.687
2025-04-07 02:14:56.908180: 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:1743992096.910303 3863697 buffer_comparator.cc:156] Difference at 35: 8.74559, expected 10.6612
E0000 00:00:1743992096.910314 3863697 buffer_comparator.cc:156] Difference at 37: 8.80435, expected 10.8965
E0000 00:00:1743992096.910317 3863697 buffer_comparator.cc:156] Difference at 41: 8.5867, expected 7.30318
E0000 00:00:1743992096.910320 3863697 buffer_comparator.cc:156] Difference at 42: 6.60885, expected 9.79879
E0000 00:00:1743992096.910323 3863697 buffer_comparator.cc:156] Difference at 43: 8.36889, expected 9.55385
E0000 00:00:1743992096.910326 3863697 buffer_comparator.cc:156] Difference at 46: 7.80689, expected 9.75357
E0000 00:00:1743992096.910329 3863697 buffer_comparator.cc:156] Difference at 47: 8.61022, expected 9.77022
E0000 00:00:1743992096.910331 3863697 buffer_comparator.cc:156] Difference at 48: 8.38568, expected 9.70077
E0000 00:00:1743992096.910334 3863697 buffer_comparator.cc:156] Difference at 50: 9.92471, expected 11.4143
E0000 00:00:1743992096.910337 3863697 buffer_comparator.cc:156] Difference at 51: 7.97977, expected 10.687
2025-04-07 02:14:56.910341: 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:1743992096.912451 3863697 buffer_comparator.cc:156] Difference at 35: 8.74559, expected 10.6612
E0000 00:00:1743992096.912462 3863697 buffer_comparator.cc:156] Difference at 37: 8.80435, expected 10.8965
E0000 00:00:1743992096.912465 3863697 buffer_comparator.cc:156] Difference at 41: 8.5867, expected 7.30318
E0000 00:00:1743992096.912468 3863697 buffer_comparator.cc:156] Difference at 42: 6.60885, expected 9.79879
E0000 00:00:1743992096.912470 3863697 buffer_comparator.cc:156] Difference at 43: 8.36889, expected 9.55385
E0000 00:00:1743992096.912473 3863697 buffer_comparator.cc:156] Difference at 46: 7.80689, expected 9.75357
E0000 00:00:1743992096.912476 3863697 buffer_comparator.cc:156] Difference at 47: 8.61022, expected 9.77022
E0000 00:00:1743992096.912479 3863697 buffer_comparator.cc:156] Difference at 48: 8.38568, expected 9.70077
E0000 00:00:1743992096.912482 3863697 buffer_comparator.cc:156] Difference at 50: 9.92471, expected 11.4143
E0000 00:00:1743992096.912485 3863697 buffer_comparator.cc:156] Difference at 51: 7.97977, expected 10.687
2025-04-07 02:14:56.912489: 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:1743992096.923768 3863697 buffer_comparator.cc:156] Difference at 1024: 0, expected 0.527629
E0000 00:00:1743992096.923780 3863697 buffer_comparator.cc:156] Difference at 1025: 0, expected 0.544573
E0000 00:00:1743992096.923783 3863697 buffer_comparator.cc:156] Difference at 1026: 0, expected 0.31953
E0000 00:00:1743992096.923786 3863697 buffer_comparator.cc:156] Difference at 1027: 0, expected 0.61527
E0000 00:00:1743992096.923789 3863697 buffer_comparator.cc:156] Difference at 1028: 0, expected 1.06241
E0000 00:00:1743992096.923792 3863697 buffer_comparator.cc:156] Difference at 1029: 0, expected 1.22093
E0000 00:00:1743992096.923795 3863697 buffer_comparator.cc:156] Difference at 1030: 0, expected 0.553754
E0000 00:00:1743992096.923797 3863697 buffer_comparator.cc:156] Difference at 1031: 0, expected 0.665092
E0000 00:00:1743992096.923800 3863697 buffer_comparator.cc:156] Difference at 1032: 0, expected 1.13956
E0000 00:00:1743992096.923803 3863697 buffer_comparator.cc:156] Difference at 1033: 0, expected 0.481583
2025-04-07 02:14:56.923810: 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:1743992096.926572 3863697 buffer_comparator.cc:156] Difference at 16: -nan, expected 7.32342
E0000 00:00:1743992096.926583 3863697 buffer_comparator.cc:156] Difference at 17: -nan, expected 10.2674
E0000 00:00:1743992096.926587 3863697 buffer_comparator.cc:156] Difference at 18: -nan, expected 8.76472
E0000 00:00:1743992096.926589 3863697 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.04411
E0000 00:00:1743992096.926592 3863697 buffer_comparator.cc:156] Difference at 20: -nan, expected 7.44869
E0000 00:00:1743992096.926595 3863697 buffer_comparator.cc:156] Difference at 21: -nan, expected 8.01524
E0000 00:00:1743992096.926597 3863697 buffer_comparator.cc:156] Difference at 22: -nan, expected 10.2221
E0000 00:00:1743992096.926600 3863697 buffer_comparator.cc:156] Difference at 23: -nan, expected 7.96118
E0000 00:00:1743992096.926603 3863697 buffer_comparator.cc:156] Difference at 24: -nan, expected 8.62185
E0000 00:00:1743992096.926605 3863697 buffer_comparator.cc:156] Difference at 25: -nan, expected 8.94114
2025-04-07 02:14:56.926610: 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:1743992096.928726 3863697 buffer_comparator.cc:156] Difference at 16: -nan, expected 7.32342
E0000 00:00:1743992096.928737 3863697 buffer_comparator.cc:156] Difference at 17: -nan, expected 10.2674
E0000 00:00:1743992096.928740 3863697 buffer_comparator.cc:156] Difference at 18: -nan, expected 8.76472
E0000 00:00:1743992096.928743 3863697 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.04411
E0000 00:00:1743992096.928745 3863697 buffer_comparator.cc:156] Difference at 20: -nan, expected 7.44869
E0000 00:00:1743992096.928748 3863697 buffer_comparator.cc:156] Difference at 21: -nan, expected 8.01524
E0000 00:00:1743992096.928751 3863697 buffer_comparator.cc:156] Difference at 22: -nan, expected 10.2221
E0000 00:00:1743992096.928753 3863697 buffer_comparator.cc:156] Difference at 23: -nan, expected 7.96118
E0000 00:00:1743992096.928756 3863697 buffer_comparator.cc:156] Difference at 24: -nan, expected 8.62185
E0000 00:00:1743992096.928759 3863697 buffer_comparator.cc:156] Difference at 25: -nan, expected 8.94114
2025-04-07 02:14:56.928763: 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:1743992096.930883 3863697 buffer_comparator.cc:156] Difference at 16: -nan, expected 7.32342
E0000 00:00:1743992096.930894 3863697 buffer_comparator.cc:156] Difference at 17: -nan, expected 10.2674
E0000 00:00:1743992096.930897 3863697 buffer_comparator.cc:156] Difference at 18: -nan, expected 8.76472
E0000 00:00:1743992096.930900 3863697 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.04411
E0000 00:00:1743992096.930903 3863697 buffer_comparator.cc:156] Difference at 20: -nan, expected 7.44869
E0000 00:00:1743992096.930905 3863697 buffer_comparator.cc:156] Difference at 21: -nan, expected 8.01524
E0000 00:00:1743992096.930908 3863697 buffer_comparator.cc:156] Difference at 22: -nan, expected 10.2221
E0000 00:00:1743992096.930911 3863697 buffer_comparator.cc:156] Difference at 23: -nan, expected 7.96118
E0000 00:00:1743992096.930913 3863697 buffer_comparator.cc:156] Difference at 24: -nan, expected 8.62185
E0000 00:00:1743992096.930916 3863697 buffer_comparator.cc:156] Difference at 25: -nan, expected 8.94114
2025-04-07 02:14:56.930920: 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:1743992096.933040 3863697 buffer_comparator.cc:156] Difference at 32: -nan, expected 7.79614
E0000 00:00:1743992096.933052 3863697 buffer_comparator.cc:156] Difference at 33: -nan, expected 8.40799
E0000 00:00:1743992096.933055 3863697 buffer_comparator.cc:156] Difference at 34: -nan, expected 7.73293
E0000 00:00:1743992096.933057 3863697 buffer_comparator.cc:156] Difference at 35: -nan, expected 9.69584
E0000 00:00:1743992096.933060 3863697 buffer_comparator.cc:156] Difference at 36: -nan, expected 8.80404
E0000 00:00:1743992096.933062 3863697 buffer_comparator.cc:156] Difference at 37: -nan, expected 10.1025
E0000 00:00:1743992096.933065 3863697 buffer_comparator.cc:156] Difference at 38: -nan, expected 8.25699
E0000 00:00:1743992096.933068 3863697 buffer_comparator.cc:156] Difference at 39: -nan, expected 8.00626
E0000 00:00:1743992096.933071 3863697 buffer_comparator.cc:156] Difference at 40: -nan, expected 8.86599
E0000 00:00:1743992096.933073 3863697 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.81407
2025-04-07 02:14:56.933077: 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:1743992096.935204 3863697 buffer_comparator.cc:156] Difference at 32: -nan, expected 7.79614
E0000 00:00:1743992096.935214 3863697 buffer_comparator.cc:156] Difference at 33: -nan, expected 8.40799
E0000 00:00:1743992096.935217 3863697 buffer_comparator.cc:156] Difference at 34: -nan, expected 7.73293
E0000 00:00:1743992096.935220 3863697 buffer_comparator.cc:156] Difference at 35: -nan, expected 9.69584
E0000 00:00:1743992096.935223 3863697 buffer_comparator.cc:156] Difference at 36: -nan, expected 8.80404
E0000 00:00:1743992096.935225 3863697 buffer_comparator.cc:156] Difference at 37: -nan, expected 10.1025
E0000 00:00:1743992096.935228 3863697 buffer_comparator.cc:156] Difference at 38: -nan, expected 8.25699
E0000 00:00:1743992096.935231 3863697 buffer_comparator.cc:156] Difference at 39: -nan, expected 8.00626
E0000 00:00:1743992096.935234 3863697 buffer_comparator.cc:156] Difference at 40: -nan, expected 8.86599
E0000 00:00:1743992096.935236 3863697 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.81407
2025-04-07 02:14:56.935240: 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:1743992096.937356 3863697 buffer_comparator.cc:156] Difference at 32: -nan, expected 7.79614
E0000 00:00:1743992096.937366 3863697 buffer_comparator.cc:156] Difference at 33: -nan, expected 8.40799
E0000 00:00:1743992096.937369 3863697 buffer_comparator.cc:156] Difference at 34: -nan, expected 7.73293
E0000 00:00:1743992096.937372 3863697 buffer_comparator.cc:156] Difference at 35: -nan, expected 9.69584
E0000 00:00:1743992096.937374 3863697 buffer_comparator.cc:156] Difference at 36: -nan, expected 8.80404
E0000 00:00:1743992096.937377 3863697 buffer_comparator.cc:156] Difference at 37: -nan, expected 10.1025
E0000 00:00:1743992096.937380 3863697 buffer_comparator.cc:156] Difference at 38: -nan, expected 8.25699
E0000 00:00:1743992096.937383 3863697 buffer_comparator.cc:156] Difference at 39: -nan, expected 8.00626
E0000 00:00:1743992096.937385 3863697 buffer_comparator.cc:156] Difference at 40: -nan, expected 8.86599
E0000 00:00:1743992096.937388 3863697 buffer_comparator.cc:156] Difference at 41: -nan, expected 7.81407
2025-04-07 02:14:56.937392: 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:1743992096.956568 3863697 buffer_comparator.cc:156] Difference at 16: 6.71717, expected 7.73985
E0000 00:00:1743992096.956579 3863697 buffer_comparator.cc:156] Difference at 17: 10.3486, expected 7.11199
E0000 00:00:1743992096.956583 3863697 buffer_comparator.cc:156] Difference at 21: 7.63301, expected 9.39119
E0000 00:00:1743992096.956586 3863697 buffer_comparator.cc:156] Difference at 25: 8.86168, expected 7.82992
E0000 00:00:1743992096.956590 3863697 buffer_comparator.cc:156] Difference at 30: 8.40614, expected 9.51389
E0000 00:00:1743992096.956593 3863697 buffer_comparator.cc:156] Difference at 32: 7.3089, expected 9.09096
E0000 00:00:1743992096.956596 3863697 buffer_comparator.cc:156] Difference at 36: 7.57855, expected 8.75347
E0000 00:00:1743992096.956599 3863697 buffer_comparator.cc:156] Difference at 37: 8.77813, expected 7.34985
E0000 00:00:1743992096.956602 3863697 buffer_comparator.cc:156] Difference at 40: 8.03792, expected 6.04756
E0000 00:00:1743992096.956604 3863697 buffer_comparator.cc:156] Difference at 45: 9.4609, expected 8.12358
2025-04-07 02:14:56.956609: 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:1743992096.958723 3863697 buffer_comparator.cc:156] Difference at 16: 6.71717, expected 7.73985
E0000 00:00:1743992096.958733 3863697 buffer_comparator.cc:156] Difference at 17: 10.3486, expected 7.11199
E0000 00:00:1743992096.958737 3863697 buffer_comparator.cc:156] Difference at 21: 7.63301, expected 9.39119
E0000 00:00:1743992096.958740 3863697 buffer_comparator.cc:156] Difference at 25: 8.86168, expected 7.82992
E0000 00:00:1743992096.958743 3863697 buffer_comparator.cc:156] Difference at 30: 8.40614, expected 9.51389
E0000 00:00:1743992096.958745 3863697 buffer_comparator.cc:156] Difference at 32: 7.3089, expected 9.09096
E0000 00:00:1743992096.958748 3863697 buffer_comparator.cc:156] Difference at 36: 7.57855, expected 8.75347
E0000 00:00:1743992096.958751 3863697 buffer_comparator.cc:156] Difference at 37: 8.77813, expected 7.34985
E0000 00:00:1743992096.958754 3863697 buffer_comparator.cc:156] Difference at 40: 8.03792, expected 6.04756
E0000 00:00:1743992096.958757 3863697 buffer_comparator.cc:156] Difference at 45: 9.4609, expected 8.12358
2025-04-07 02:14:56.958761: 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:1743992096.960876 3863697 buffer_comparator.cc:156] Difference at 32: 7.3089, expected 9.09096
E0000 00:00:1743992096.960887 3863697 buffer_comparator.cc:156] Difference at 36: 7.57855, expected 8.75347
E0000 00:00:1743992096.960890 3863697 buffer_comparator.cc:156] Difference at 37: 8.77813, expected 7.34985
E0000 00:00:1743992096.960892 3863697 buffer_comparator.cc:156] Difference at 40: 8.03792, expected 6.04756
E0000 00:00:1743992096.960895 3863697 buffer_comparator.cc:156] Difference at 45: 9.4609, expected 8.12358
E0000 00:00:1743992096.960898 3863697 buffer_comparator.cc:156] Difference at 46: 9.46794, expected 7.84264
E0000 00:00:1743992096.960901 3863697 buffer_comparator.cc:156] Difference at 48: 9.57179, expected 7.26336
E0000 00:00:1743992096.960904 3863697 buffer_comparator.cc:156] Difference at 51: 9.54877, expected 6.79213
E0000 00:00:1743992096.960906 3863697 buffer_comparator.cc:156] Difference at 52: 8.7149, expected 6.91881
E0000 00:00:1743992096.960909 3863697 buffer_comparator.cc:156] Difference at 53: 8.27737, expected 9.47379
2025-04-07 02:14:56.960914: 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:1743992096.963022 3863697 buffer_comparator.cc:156] Difference at 32: 7.3089, expected 9.09096
E0000 00:00:1743992096.963033 3863697 buffer_comparator.cc:156] Difference at 36: 7.57855, expected 8.75347
E0000 00:00:1743992096.963036 3863697 buffer_comparator.cc:156] Difference at 37: 8.77813, expected 7.34985
E0000 00:00:1743992096.963039 3863697 buffer_comparator.cc:156] Difference at 40: 8.03792, expected 6.04756
E0000 00:00:1743992096.963042 3863697 buffer_comparator.cc:156] Difference at 45: 9.4609, expected 8.12358
E0000 00:00:1743992096.963045 3863697 buffer_comparator.cc:156] Difference at 46: 9.46794, expected 7.84264
E0000 00:00:1743992096.963050 3863697 buffer_comparator.cc:156] Difference at 48: 9.57179, expected 7.26336
E0000 00:00:1743992096.963053 3863697 buffer_comparator.cc:156] Difference at 51: 9.54877, expected 6.79213
E0000 00:00:1743992096.963055 3863697 buffer_comparator.cc:156] Difference at 52: 8.7149, expected 6.91881
E0000 00:00:1743992096.963058 3863697 buffer_comparator.cc:156] Difference at 53: 8.27737, expected 9.47379
2025-04-07 02:14:56.963063: 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:1743992096.965182 3863697 buffer_comparator.cc:156] Difference at 64: 6.85353, expected 8.43679
E0000 00:00:1743992096.965192 3863697 buffer_comparator.cc:156] Difference at 65: 7.78097, expected 9.73984
E0000 00:00:1743992096.965195 3863697 buffer_comparator.cc:156] Difference at 66: 8.16401, expected 10.0838
E0000 00:00:1743992096.965198 3863697 buffer_comparator.cc:156] Difference at 68: 7.2158, expected 9.00042
E0000 00:00:1743992096.965201 3863697 buffer_comparator.cc:156] Difference at 72: 7.39918, expected 8.47292
E0000 00:00:1743992096.965204 3863697 buffer_comparator.cc:156] Difference at 73: 7.1204, expected 8.09814
E0000 00:00:1743992096.965207 3863697 buffer_comparator.cc:156] Difference at 74: 6.09996, expected 8.89678
E0000 00:00:1743992096.965210 3863697 buffer_comparator.cc:156] Difference at 75: 6.72143, expected 9.07456
E0000 00:00:1743992096.965212 3863697 buffer_comparator.cc:156] Difference at 76: 6.65785, expected 8.9659
E0000 00:00:1743992096.965215 3863697 buffer_comparator.cc:156] Difference at 77: 6.49377, expected 9.28323
2025-04-07 02:14:56.965220: 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:1743992096.967337 3863697 buffer_comparator.cc:156] Difference at 64: 6.85353, expected 8.43679
E0000 00:00:1743992096.967347 3863697 buffer_comparator.cc:156] Difference at 65: 7.78097, expected 9.73984
E0000 00:00:1743992096.967350 3863697 buffer_comparator.cc:156] Difference at 66: 8.16401, expected 10.0838
E0000 00:00:1743992096.967354 3863697 buffer_comparator.cc:156] Difference at 68: 7.2158, expected 9.00042
E0000 00:00:1743992096.967357 3863697 buffer_comparator.cc:156] Difference at 72: 7.39918, expected 8.47292
E0000 00:00:1743992096.967359 3863697 buffer_comparator.cc:156] Difference at 73: 7.1204, expected 8.09814
E0000 00:00:1743992096.967362 3863697 buffer_comparator.cc:156] Difference at 74: 6.09996, expected 8.89678
E0000 00:00:1743992096.967365 3863697 buffer_comparator.cc:156] Difference at 75: 6.72143, expected 9.07456
E0000 00:00:1743992096.967367 3863697 buffer_comparator.cc:156] Difference at 76: 6.65785, expected 8.9659
E0000 00:00:1743992096.967370 3863697 buffer_comparator.cc:156] Difference at 77: 6.49377, expected 9.28323
2025-04-07 02:14:56.967375: 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:1743992096.969494 3863697 buffer_comparator.cc:156] Difference at 64: 6.85353, expected 8.43679
E0000 00:00:1743992096.969505 3863697 buffer_comparator.cc:156] Difference at 65: 7.78097, expected 9.73984
E0000 00:00:1743992096.969508 3863697 buffer_comparator.cc:156] Difference at 66: 8.16401, expected 10.0838
E0000 00:00:1743992096.969511 3863697 buffer_comparator.cc:156] Difference at 68: 7.2158, expected 9.00042
E0000 00:00:1743992096.969514 3863697 buffer_comparator.cc:156] Difference at 72: 7.39918, expected 8.47292
E0000 00:00:1743992096.969516 3863697 buffer_comparator.cc:156] Difference at 73: 7.1204, expected 8.09814
E0000 00:00:1743992096.969519 3863697 buffer_comparator.cc:156] Difference at 74: 6.09996, expected 8.89678
E0000 00:00:1743992096.969522 3863697 buffer_comparator.cc:156] Difference at 75: 6.72143, expected 9.07456
E0000 00:00:1743992096.969527 3863697 buffer_comparator.cc:156] Difference at 76: 6.65785, expected 8.9659
E0000 00:00:1743992096.969531 3863697 buffer_comparator.cc:156] Difference at 77: 6.49377, expected 9.28323
2025-04-07 02:14:56.969536: 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:1743992096.971652 3863697 buffer_comparator.cc:156] Difference at 128: 7.99514, expected 10.3054
E0000 00:00:1743992096.971663 3863697 buffer_comparator.cc:156] Difference at 129: 8.05157, expected 9.87733
E0000 00:00:1743992096.971666 3863697 buffer_comparator.cc:156] Difference at 132: 6.83475, expected 9.08348
E0000 00:00:1743992096.971669 3863697 buffer_comparator.cc:156] Difference at 133: 8.63235, expected 9.79527
E0000 00:00:1743992096.971672 3863697 buffer_comparator.cc:156] Difference at 134: 9.3024, expected 7.65675
E0000 00:00:1743992096.971675 3863697 buffer_comparator.cc:156] Difference at 137: 6.81486, expected 8.67269
E0000 00:00:1743992096.971677 3863697 buffer_comparator.cc:156] Difference at 145: 10.3404, expected 7.81371
E0000 00:00:1743992096.971680 3863697 buffer_comparator.cc:156] Difference at 147: 7.11057, expected 9.22175
E0000 00:00:1743992096.971683 3863697 buffer_comparator.cc:156] Difference at 150: 9.55365, expected 8.16186
E0000 00:00:1743992096.971686 3863697 buffer_comparator.cc:156] Difference at 153: 9.27692, expected 10.579
2025-04-07 02:14:56.971691: 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:1743992096.973805 3863697 buffer_comparator.cc:156] Difference at 128: 7.99514, expected 10.3054
E0000 00:00:1743992096.973816 3863697 buffer_comparator.cc:156] Difference at 129: 8.05157, expected 9.87733
E0000 00:00:1743992096.973819 3863697 buffer_comparator.cc:156] Difference at 132: 6.83475, expected 9.08348
E0000 00:00:1743992096.973822 3863697 buffer_comparator.cc:156] Difference at 133: 8.63235, expected 9.79527
E0000 00:00:1743992096.973825 3863697 buffer_comparator.cc:156] Difference at 134: 9.3024, expected 7.65675
E0000 00:00:1743992096.973828 3863697 buffer_comparator.cc:156] Difference at 137: 6.81486, expected 8.67269
E0000 00:00:1743992096.973831 3863697 buffer_comparator.cc:156] Difference at 145: 10.3404, expected 7.81371
E0000 00:00:1743992096.973834 3863697 buffer_comparator.cc:156] Difference at 147: 7.11057, expected 9.22175
E0000 00:00:1743992096.973836 3863697 buffer_comparator.cc:156] Difference at 150: 9.55365, expected 8.16186
E0000 00:00:1743992096.973839 3863697 buffer_comparator.cc:156] Difference at 153: 9.27692, expected 10.579
2025-04-07 02:14:56.973844: 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:1743992096.981839 3863697 buffer_comparator.cc:156] Difference at 16: 0, expected 11.6059
E0000 00:00:1743992096.981851 3863697 buffer_comparator.cc:156] Difference at 17: 0, expected 14.502
E0000 00:00:1743992096.981854 3863697 buffer_comparator.cc:156] Difference at 18: 0, expected 11.2449
E0000 00:00:1743992096.981857 3863697 buffer_comparator.cc:156] Difference at 19: 0, expected 10.0998
E0000 00:00:1743992096.981860 3863697 buffer_comparator.cc:156] Difference at 20: 0, expected 14.0222
E0000 00:00:1743992096.981863 3863697 buffer_comparator.cc:156] Difference at 21: 0, expected 10.1321
E0000 00:00:1743992096.981866 3863697 buffer_comparator.cc:156] Difference at 22: 0, expected 10.2986
E0000 00:00:1743992096.981869 3863697 buffer_comparator.cc:156] Difference at 23: 0, expected 14.1109
E0000 00:00:1743992096.981872 3863697 buffer_comparator.cc:156] Difference at 24: 0, expected 13.3463
E0000 00:00:1743992096.981875 3863697 buffer_comparator.cc:156] Difference at 25: 0, expected 12.8369
2025-04-07 02:14:56.981880: 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:1743992096.984813 3863697 buffer_comparator.cc:156] Difference at 16: 0, expected 11.6059
E0000 00:00:1743992096.984824 3863697 buffer_comparator.cc:156] Difference at 17: 0, expected 14.502
E0000 00:00:1743992096.984827 3863697 buffer_comparator.cc:156] Difference at 18: 0, expected 11.2449
E0000 00:00:1743992096.984830 3863697 buffer_comparator.cc:156] Difference at 19: 0, expected 10.0998
E0000 00:00:1743992096.984833 3863697 buffer_comparator.cc:156] Difference at 20: 0, expected 14.0222
E0000 00:00:1743992096.984836 3863697 buffer_comparator.cc:156] Difference at 21: 0, expected 10.1321
E0000 00:00:1743992096.984839 3863697 buffer_comparator.cc:156] Difference at 22: 0, expected 10.2986
E0000 00:00:1743992096.984842 3863697 buffer_comparator.cc:156] Difference at 23: 0, expected 14.1109
E0000 00:00:1743992096.984845 3863697 buffer_comparator.cc:156] Difference at 24: 0, expected 13.3463
E0000 00:00:1743992096.984848 3863697 buffer_comparator.cc:156] Difference at 25: 0, expected 12.8369
2025-04-07 02:14:56.984852: 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:1743992096.987791 3863697 buffer_comparator.cc:156] Difference at 16: 0, expected 11.6059
E0000 00:00:1743992096.987803 3863697 buffer_comparator.cc:156] Difference at 17: 0, expected 14.502
E0000 00:00:1743992096.987806 3863697 buffer_comparator.cc:156] Difference at 18: 0, expected 11.2449
E0000 00:00:1743992096.987809 3863697 buffer_comparator.cc:156] Difference at 19: 0, expected 10.0998
E0000 00:00:1743992096.987812 3863697 buffer_comparator.cc:156] Difference at 20: 0, expected 14.0222
E0000 00:00:1743992096.987815 3863697 buffer_comparator.cc:156] Difference at 21: 0, expected 10.1321
E0000 00:00:1743992096.987818 3863697 buffer_comparator.cc:156] Difference at 22: 0, expected 10.2986
E0000 00:00:1743992096.987821 3863697 buffer_comparator.cc:156] Difference at 23: 0, expected 14.1109
E0000 00:00:1743992096.987823 3863697 buffer_comparator.cc:156] Difference at 24: 0, expected 13.3463
E0000 00:00:1743992096.987826 3863697 buffer_comparator.cc:156] Difference at 25: 0, expected 12.8369
2025-04-07 02:14:56.987831: 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:1743992096.990312 3863697 buffer_comparator.cc:156] Difference at 32: 0, expected 12.4
E0000 00:00:1743992096.990323 3863697 buffer_comparator.cc:156] Difference at 33: 0, expected 12.9454
E0000 00:00:1743992096.990326 3863697 buffer_comparator.cc:156] Difference at 34: 0, expected 12.9462
E0000 00:00:1743992096.990329 3863697 buffer_comparator.cc:156] Difference at 35: 0, expected 13.9775
E0000 00:00:1743992096.990332 3863697 buffer_comparator.cc:156] Difference at 36: 0, expected 15.0433
E0000 00:00:1743992096.990335 3863697 buffer_comparator.cc:156] Difference at 37: 0, expected 12.0589
E0000 00:00:1743992096.990338 3863697 buffer_comparator.cc:156] Difference at 38: 0, expected 14.4629
E0000 00:00:1743992096.990341 3863697 buffer_comparator.cc:156] Difference at 39: 0, expected 12.7671
E0000 00:00:1743992096.990344 3863697 buffer_comparator.cc:156] Difference at 40: 0, expected 12.3584
E0000 00:00:1743992096.990347 3863697 buffer_comparator.cc:156] Difference at 41: 0, expected 11.6002
2025-04-07 02:14:56.990351: 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:1743992096.992813 3863697 buffer_comparator.cc:156] Difference at 32: 0, expected 12.4
E0000 00:00:1743992096.992824 3863697 buffer_comparator.cc:156] Difference at 33: 0, expected 12.9454
E0000 00:00:1743992096.992827 3863697 buffer_comparator.cc:156] Difference at 34: 0, expected 12.9462
E0000 00:00:1743992096.992832 3863697 buffer_comparator.cc:156] Difference at 35: 0, expected 13.9775
E0000 00:00:1743992096.992834 3863697 buffer_comparator.cc:156] Difference at 36: 0, expected 15.0433
E0000 00:00:1743992096.992837 3863697 buffer_comparator.cc:156] Difference at 37: 0, expected 12.0589
E0000 00:00:1743992096.992840 3863697 buffer_comparator.cc:156] Difference at 38: 0, expected 14.4629
E0000 00:00:1743992096.992843 3863697 buffer_comparator.cc:156] Difference at 39: 0, expected 12.7671
E0000 00:00:1743992096.992846 3863697 buffer_comparator.cc:156] Difference at 40: 0, expected 12.3584
E0000 00:00:1743992096.992849 3863697 buffer_comparator.cc:156] Difference at 41: 0, expected 11.6002
2025-04-07 02:14:56.992853: 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:1743992096.995305 3863697 buffer_comparator.cc:156] Difference at 32: 0, expected 12.4
E0000 00:00:1743992096.995316 3863697 buffer_comparator.cc:156] Difference at 33: 0, expected 12.9454
E0000 00:00:1743992096.995319 3863697 buffer_comparator.cc:156] Difference at 34: 0, expected 12.9462
E0000 00:00:1743992096.995322 3863697 buffer_comparator.cc:156] Difference at 35: 0, expected 13.9775
E0000 00:00:1743992096.995325 3863697 buffer_comparator.cc:156] Difference at 36: 0, expected 15.0433
E0000 00:00:1743992096.995328 3863697 buffer_comparator.cc:156] Difference at 37: 0, expected 12.0589
E0000 00:00:1743992096.995331 3863697 buffer_comparator.cc:156] Difference at 38: 0, expected 14.4629
E0000 00:00:1743992096.995334 3863697 buffer_comparator.cc:156] Difference at 39: 0, expected 12.7671
E0000 00:00:1743992096.995337 3863697 buffer_comparator.cc:156] Difference at 40: 0, expected 12.3584
E0000 00:00:1743992096.995340 3863697 buffer_comparator.cc:156] Difference at 41: 0, expected 11.6002
2025-04-07 02:14:56.995344: 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:1743992097.019931 3863697 buffer_comparator.cc:156] Difference at 16: -nan, expected 9.17386
E0000 00:00:1743992097.019942 3863697 buffer_comparator.cc:156] Difference at 17: -nan, expected 8.63962
E0000 00:00:1743992097.019945 3863697 buffer_comparator.cc:156] Difference at 18: -nan, expected 11.0889
E0000 00:00:1743992097.019948 3863697 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.96333
E0000 00:00:1743992097.019951 3863697 buffer_comparator.cc:156] Difference at 20: -nan, expected 8.01067
E0000 00:00:1743992097.019954 3863697 buffer_comparator.cc:156] Difference at 21: -nan, expected 11.0277
E0000 00:00:1743992097.019957 3863697 buffer_comparator.cc:156] Difference at 22: -nan, expected 7.61027
E0000 00:00:1743992097.019959 3863697 buffer_comparator.cc:156] Difference at 23: -nan, expected 7.70495
E0000 00:00:1743992097.019962 3863697 buffer_comparator.cc:156] Difference at 24: -nan, expected 10.9025
E0000 00:00:1743992097.019965 3863697 buffer_comparator.cc:156] Difference at 25: -nan, expected 9.9377
2025-04-07 02:14:57.019969: 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:1743992097.021765 3863697 buffer_comparator.cc:156] Difference at 16: -nan, expected 9.17386
E0000 00:00:1743992097.021775 3863697 buffer_comparator.cc:156] Difference at 17: -nan, expected 8.63962
E0000 00:00:1743992097.021778 3863697 buffer_comparator.cc:156] Difference at 18: -nan, expected 11.0889
E0000 00:00:1743992097.021781 3863697 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.96333
E0000 00:00:1743992097.021784 3863697 buffer_comparator.cc:156] Difference at 20: -nan, expected 8.01067
E0000 00:00:1743992097.021786 3863697 buffer_comparator.cc:156] Difference at 21: -nan, expected 11.0277
E0000 00:00:1743992097.021789 3863697 buffer_comparator.cc:156] Difference at 22: -nan, expected 7.61027
E0000 00:00:1743992097.021794 3863697 buffer_comparator.cc:156] Difference at 23: -nan, expected 7.70495
E0000 00:00:1743992097.021797 3863697 buffer_comparator.cc:156] Difference at 24: -nan, expected 10.9025
E0000 00:00:1743992097.021799 3863697 buffer_comparator.cc:156] Difference at 25: -nan, expected 9.9377
2025-04-07 02:14:57.021804: 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:1743992097.023612 3863697 buffer_comparator.cc:156] Difference at 16: -nan, expected 9.17386
E0000 00:00:1743992097.023622 3863697 buffer_comparator.cc:156] Difference at 17: -nan, expected 8.63962
E0000 00:00:1743992097.023625 3863697 buffer_comparator.cc:156] Difference at 18: -nan, expected 11.0889
E0000 00:00:1743992097.023628 3863697 buffer_comparator.cc:156] Difference at 19: -nan, expected 8.96333
E0000 00:00:1743992097.023631 3863697 buffer_comparator.cc:156] Difference at 20: -nan, expected 8.01067
E0000 00:00:1743992097.023634 3863697 buffer_comparator.cc:156] Difference at 21: -nan, expected 11.0277
E0000 00:00:1743992097.023637 3863697 buffer_comparator.cc:156] Difference at 22: -nan, expected 7.61027
E0000 00:00:1743992097.023639 3863697 buffer_comparator.cc:156] Difference at 23: -nan, expected 7.70495
E0000 00:00:1743992097.023642 3863697 buffer_comparator.cc:156] Difference at 24: -nan, expected 10.9025
E0000 00:00:1743992097.023645 3863697 buffer_comparator.cc:156] Difference at 25: -nan, expected 9.9377
2025-04-07 02:14:57.023649: 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:1743992097.025441 3863697 buffer_comparator.cc:156] Difference at 32: -nan, expected 9.26776
E0000 00:00:1743992097.025451 3863697 buffer_comparator.cc:156] Difference at 33: -nan, expected 9.1446
E0000 00:00:1743992097.025454 3863697 buffer_comparator.cc:156] Difference at 34: -nan, expected 9.79374
E0000 00:00:1743992097.025457 3863697 buffer_comparator.cc:156] Difference at 35: -nan, expected 10.23
E0000 00:00:1743992097.025460 3863697 buffer_comparator.cc:156] Difference at 36: -nan, expected 10.448
E0000 00:00:1743992097.025463 3863697 buffer_comparator.cc:156] Difference at 37: -nan, expected 11.8017
E0000 00:00:1743992097.025466 3863697 buffer_comparator.cc:156] Difference at 38: -nan, expected 9.69309
E0000 00:00:1743992097.025468 3863697 buffer_comparator.cc:156] Difference at 39: -nan, expected 11.2277
E0000 00:00:1743992097.025471 3863697 buffer_comparator.cc:156] Difference at 40: -nan, expected 10.0359
E0000 00:00:1743992097.025474 3863697 buffer_comparator.cc:156] Difference at 41: -nan, expected 9.30719
2025-04-07 02:14:57.025478: 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:1743992097.027275 3863697 buffer_comparator.cc:156] Difference at 32: -nan, expected 9.26776
E0000 00:00:1743992097.027286 3863697 buffer_comparator.cc:156] Difference at 33: -nan, expected 9.1446
E0000 00:00:1743992097.027289 3863697 buffer_comparator.cc:156] Difference at 34: -nan, expected 9.79374
E0000 00:00:1743992097.027292 3863697 buffer_comparator.cc:156] Difference at 35: -nan, expected 10.23
E0000 00:00:1743992097.027295 3863697 buffer_comparator.cc:156] Difference at 36: -nan, expected 10.448
E0000 00:00:1743992097.027298 3863697 buffer_comparator.cc:156] Difference at 37: -nan, expected 11.8017
E0000 00:00:1743992097.027300 3863697 buffer_comparator.cc:156] Difference at 38: -nan, expected 9.69309
E0000 00:00:1743992097.027303 3863697 buffer_comparator.cc:156] Difference at 39: -nan, expected 11.2277
E0000 00:00:1743992097.027306 3863697 buffer_comparator.cc:156] Difference at 40: -nan, expected 10.0359
E0000 00:00:1743992097.027309 3863697 buffer_comparator.cc:156] Difference at 41: -nan, expected 9.30719
2025-04-07 02:14:57.027315: 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:1743992097.029106 3863697 buffer_comparator.cc:156] Difference at 32: -nan, expected 9.26776
E0000 00:00:1743992097.029117 3863697 buffer_comparator.cc:156] Difference at 33: -nan, expected 9.1446
E0000 00:00:1743992097.029120 3863697 buffer_comparator.cc:156] Difference at 34: -nan, expected 9.79374
E0000 00:00:1743992097.029122 3863697 buffer_comparator.cc:156] Difference at 35: -nan, expected 10.23
E0000 00:00:1743992097.029125 3863697 buffer_comparator.cc:156] Difference at 36: -nan, expected 10.448
E0000 00:00:1743992097.029128 3863697 buffer_comparator.cc:156] Difference at 37: -nan, expected 11.8017
E0000 00:00:1743992097.029131 3863697 buffer_comparator.cc:156] Difference at 38: -nan, expected 9.69309
E0000 00:00:1743992097.029134 3863697 buffer_comparator.cc:156] Difference at 39: -nan, expected 11.2277
E0000 00:00:1743992097.029136 3863697 buffer_comparator.cc:156] Difference at 40: -nan, expected 10.0359
E0000 00:00:1743992097.029139 3863697 buffer_comparator.cc:156] Difference at 41: -nan, expected 9.30719
2025-04-07 02:14:57.029143: 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.158608437 (3.158608437) 	 Physics Loss: 1.981972933 (1.981972933) 	 Data Loss: 0.578243732 (0.578243732) 	 BC Loss: 0.598391712 (0.598391712)
Iteration: [  1001/ 50000] 	 Loss: 0.034501456 (0.028619381) 	 Physics Loss: 0.000539157 (0.000465246) 	 Data Loss: 0.022232454 (0.011837148) 	 BC Loss: 0.011729846 (0.016316988)
Iteration: [  2001/ 50000] 	 Loss: 0.027032075 (0.030060001) 	 Physics Loss: 0.001218293 (0.001170262) 	 Data Loss: 0.012028398 (0.011948824) 	 BC Loss: 0.013785385 (0.016940912)
Iteration: [  3001/ 50000] 	 Loss: 0.028781807 (0.025681779) 	 Physics Loss: 0.009173925 (0.004241092) 	 Data Loss: 0.004293054 (0.007728362) 	 BC Loss: 0.015314828 (0.013712320)
Iteration: [  4001/ 50000] 	 Loss: 0.049193092 (0.037200678) 	 Physics Loss: 0.002656172 (0.002551769) 	 Data Loss: 0.018408582 (0.013924235) 	 BC Loss: 0.028128337 (0.020724675)
Iteration: [  5001/ 50000] 	 Loss: 0.019937733 (0.033174727) 	 Physics Loss: 0.000668856 (0.005237338) 	 Data Loss: 0.008721712 (0.011264897) 	 BC Loss: 0.010547165 (0.016672486)
Iteration: [  6001/ 50000] 	 Loss: 0.024159394 (0.020003816) 	 Physics Loss: 0.000365201 (0.000672787) 	 Data Loss: 0.006351952 (0.006960358) 	 BC Loss: 0.017442241 (0.012370671)
Iteration: [  7001/ 50000] 	 Loss: 0.018610070 (0.021088485) 	 Physics Loss: 0.000588019 (0.000647889) 	 Data Loss: 0.007247762 (0.008391906) 	 BC Loss: 0.010774288 (0.012048689)
Iteration: [  8001/ 50000] 	 Loss: 0.020194240 (0.019458346) 	 Physics Loss: 0.000967226 (0.000871734) 	 Data Loss: 0.005540119 (0.006786535) 	 BC Loss: 0.013686894 (0.011800079)
Iteration: [  9001/ 50000] 	 Loss: 0.020544574 (0.019427346) 	 Physics Loss: 0.001345180 (0.000852766) 	 Data Loss: 0.004664297 (0.007215607) 	 BC Loss: 0.014535096 (0.011358970)
Iteration: [ 10001/ 50000] 	 Loss: 0.015292862 (0.018986428) 	 Physics Loss: 0.000520102 (0.000905526) 	 Data Loss: 0.008151696 (0.007172311) 	 BC Loss: 0.006621065 (0.010908589)
Iteration: [ 11001/ 50000] 	 Loss: 0.022855435 (0.018381279) 	 Physics Loss: 0.000386268 (0.000429100) 	 Data Loss: 0.009592065 (0.007024384) 	 BC Loss: 0.012877102 (0.010927795)
Iteration: [ 12001/ 50000] 	 Loss: 0.014141575 (0.017787997) 	 Physics Loss: 0.000836132 (0.000534793) 	 Data Loss: 0.007425712 (0.007124355) 	 BC Loss: 0.005879730 (0.010128846)
Iteration: [ 13001/ 50000] 	 Loss: 0.010630632 (0.017406996) 	 Physics Loss: 0.000509437 (0.000616333) 	 Data Loss: 0.003215889 (0.006511343) 	 BC Loss: 0.006905306 (0.010279317)
Iteration: [ 14001/ 50000] 	 Loss: 0.013067488 (0.015925726) 	 Physics Loss: 0.000660870 (0.000627965) 	 Data Loss: 0.006079779 (0.005981214) 	 BC Loss: 0.006326839 (0.009316549)
Iteration: [ 15001/ 50000] 	 Loss: 0.014089976 (0.016083684) 	 Physics Loss: 0.000765222 (0.000820126) 	 Data Loss: 0.000957327 (0.006192476) 	 BC Loss: 0.012367427 (0.009071083)
Iteration: [ 16001/ 50000] 	 Loss: 0.011554481 (0.016687429) 	 Physics Loss: 0.001439288 (0.001081884) 	 Data Loss: 0.005057607 (0.005356309) 	 BC Loss: 0.005057587 (0.010249237)
Iteration: [ 17001/ 50000] 	 Loss: 0.008298386 (0.015040044) 	 Physics Loss: 0.000735596 (0.001127110) 	 Data Loss: 0.001235073 (0.004865414) 	 BC Loss: 0.006327718 (0.009047518)
Iteration: [ 18001/ 50000] 	 Loss: 0.020617612 (0.013146948) 	 Physics Loss: 0.000790122 (0.001101069) 	 Data Loss: 0.006835818 (0.003712685) 	 BC Loss: 0.012991672 (0.008333196)
Iteration: [ 19001/ 50000] 	 Loss: 0.013813432 (0.015248597) 	 Physics Loss: 0.000885520 (0.001074361) 	 Data Loss: 0.006169923 (0.004109941) 	 BC Loss: 0.006757990 (0.010064295)
Iteration: [ 20001/ 50000] 	 Loss: 0.009095352 (0.013940509) 	 Physics Loss: 0.002372519 (0.001133401) 	 Data Loss: 0.004173467 (0.003184182) 	 BC Loss: 0.002549367 (0.009622927)
Iteration: [ 21001/ 50000] 	 Loss: 0.007158810 (0.012354066) 	 Physics Loss: 0.001433884 (0.001144025) 	 Data Loss: 0.002091464 (0.003593219) 	 BC Loss: 0.003633461 (0.007616822)
Iteration: [ 22001/ 50000] 	 Loss: 0.013130328 (0.013791627) 	 Physics Loss: 0.001378646 (0.001276990) 	 Data Loss: 0.008202977 (0.002791563) 	 BC Loss: 0.003548704 (0.009723072)
Iteration: [ 23001/ 50000] 	 Loss: 0.015937351 (0.010835836) 	 Physics Loss: 0.001402167 (0.001362659) 	 Data Loss: 0.005075668 (0.003221619) 	 BC Loss: 0.009459516 (0.006251558)
Iteration: [ 24001/ 50000] 	 Loss: 0.020060750 (0.012162323) 	 Physics Loss: 0.001117563 (0.001275578) 	 Data Loss: 0.001759502 (0.002794140) 	 BC Loss: 0.017183684 (0.008092605)
Iteration: [ 25001/ 50000] 	 Loss: 0.011357602 (0.010663286) 	 Physics Loss: 0.001268513 (0.001450315) 	 Data Loss: 0.002156476 (0.002441080) 	 BC Loss: 0.007932613 (0.006771891)
Iteration: [ 26001/ 50000] 	 Loss: 0.011544775 (0.008691516) 	 Physics Loss: 0.001099774 (0.001409120) 	 Data Loss: 0.001979375 (0.002304285) 	 BC Loss: 0.008465625 (0.004978111)
Iteration: [ 27001/ 50000] 	 Loss: 0.015457619 (0.009291158) 	 Physics Loss: 0.000908111 (0.001451046) 	 Data Loss: 0.003038743 (0.002204968) 	 BC Loss: 0.011510765 (0.005635142)
Iteration: [ 28001/ 50000] 	 Loss: 0.004877396 (0.008904044) 	 Physics Loss: 0.002144316 (0.002060717) 	 Data Loss: 0.001298780 (0.002208986) 	 BC Loss: 0.001434302 (0.004634342)
Iteration: [ 29001/ 50000] 	 Loss: 0.006905883 (0.007878463) 	 Physics Loss: 0.001174045 (0.001522218) 	 Data Loss: 0.003223591 (0.001916611) 	 BC Loss: 0.002508248 (0.004439634)
Iteration: [ 30001/ 50000] 	 Loss: 0.007669936 (0.007611540) 	 Physics Loss: 0.001393154 (0.001635356) 	 Data Loss: 0.000499706 (0.001890983) 	 BC Loss: 0.005777076 (0.004085200)
Iteration: [ 31001/ 50000] 	 Loss: 0.009445012 (0.007356300) 	 Physics Loss: 0.001389000 (0.001926766) 	 Data Loss: 0.000558271 (0.001592355) 	 BC Loss: 0.007497741 (0.003837179)
Iteration: [ 32001/ 50000] 	 Loss: 0.014761746 (0.008018765) 	 Physics Loss: 0.005334048 (0.002689240) 	 Data Loss: 0.003802153 (0.001825798) 	 BC Loss: 0.005625546 (0.003503727)
Iteration: [ 33001/ 50000] 	 Loss: 0.004137085 (0.005539955) 	 Physics Loss: 0.001168507 (0.001334399) 	 Data Loss: 0.000628497 (0.001203599) 	 BC Loss: 0.002340082 (0.003001957)
Iteration: [ 34001/ 50000] 	 Loss: 0.006806533 (0.005543606) 	 Physics Loss: 0.000993634 (0.001714336) 	 Data Loss: 0.000688867 (0.001300389) 	 BC Loss: 0.005124032 (0.002528881)
Iteration: [ 35001/ 50000] 	 Loss: 0.002602153 (0.004925680) 	 Physics Loss: 0.001339400 (0.001858316) 	 Data Loss: 0.000425479 (0.000879648) 	 BC Loss: 0.000837275 (0.002187716)
Iteration: [ 36001/ 50000] 	 Loss: 0.005823607 (0.004171041) 	 Physics Loss: 0.001609279 (0.001522730) 	 Data Loss: 0.002328272 (0.000967416) 	 BC Loss: 0.001886056 (0.001680896)
Iteration: [ 37001/ 50000] 	 Loss: 0.002417471 (0.003602090) 	 Physics Loss: 0.000783447 (0.001255809) 	 Data Loss: 0.000527038 (0.000950004) 	 BC Loss: 0.001106986 (0.001396276)
Iteration: [ 38001/ 50000] 	 Loss: 0.004167955 (0.003407350) 	 Physics Loss: 0.000959939 (0.001356680) 	 Data Loss: 0.001092754 (0.000731044) 	 BC Loss: 0.002115263 (0.001319626)
Iteration: [ 39001/ 50000] 	 Loss: 0.002879463 (0.002971070) 	 Physics Loss: 0.001146135 (0.001284969) 	 Data Loss: 0.000324087 (0.000686376) 	 BC Loss: 0.001409241 (0.000999725)
Iteration: [ 40001/ 50000] 	 Loss: 0.004722770 (0.003168148) 	 Physics Loss: 0.001417258 (0.001560614) 	 Data Loss: 0.001735250 (0.000698790) 	 BC Loss: 0.001570262 (0.000908744)
Iteration: [ 41001/ 50000] 	 Loss: 0.001900060 (0.002704850) 	 Physics Loss: 0.000999555 (0.001382467) 	 Data Loss: 0.000738261 (0.000631339) 	 BC Loss: 0.000162244 (0.000691044)
Iteration: [ 42001/ 50000] 	 Loss: 0.001598274 (0.002092211) 	 Physics Loss: 0.001005989 (0.000993158) 	 Data Loss: 0.000384238 (0.000508227) 	 BC Loss: 0.000208047 (0.000590825)
Iteration: [ 43001/ 50000] 	 Loss: 0.002319615 (0.002854182) 	 Physics Loss: 0.001340219 (0.001710652) 	 Data Loss: 0.000256991 (0.000504172) 	 BC Loss: 0.000722405 (0.000639357)
Iteration: [ 44001/ 50000] 	 Loss: 0.001560866 (0.002511591) 	 Physics Loss: 0.000632470 (0.001344989) 	 Data Loss: 0.000230765 (0.000581985) 	 BC Loss: 0.000697631 (0.000584617)
Iteration: [ 45001/ 50000] 	 Loss: 0.002259991 (0.001815259) 	 Physics Loss: 0.001150392 (0.001018210) 	 Data Loss: 0.000397562 (0.000409101) 	 BC Loss: 0.000712037 (0.000387948)
Iteration: [ 46001/ 50000] 	 Loss: 0.001558127 (0.001862627) 	 Physics Loss: 0.000627256 (0.000959221) 	 Data Loss: 0.000392324 (0.000465535) 	 BC Loss: 0.000538547 (0.000437871)
Iteration: [ 47001/ 50000] 	 Loss: 0.001834099 (0.001855874) 	 Physics Loss: 0.001413039 (0.001151622) 	 Data Loss: 0.000193402 (0.000393873) 	 BC Loss: 0.000227658 (0.000310379)
Iteration: [ 48001/ 50000] 	 Loss: 0.001248884 (0.001723553) 	 Physics Loss: 0.000700008 (0.000946964) 	 Data Loss: 0.000272257 (0.000437598) 	 BC Loss: 0.000276620 (0.000338992)
Iteration: [ 49001/ 50000] 	 Loss: 0.001509439 (0.001469041) 	 Physics Loss: 0.000762799 (0.000850576) 	 Data Loss: 0.000458921 (0.000337953) 	 BC Loss: 0.000287719 (0.000280513)

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