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-25 01:32:02.942940: I external/xla/xla/service/service.cc:152] XLA service 0x113d68a0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2025-04-25 01:32:02.942979: 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:1745544722.943797 1321573 se_gpu_pjrt_client.cc:999] Using BFC allocator.
I0000 00:00:1745544722.943911 1321573 gpu_helpers.cc:136] XLA backend allocating 3825205248 bytes on device 0 for BFCAllocator.
I0000 00:00:1745544722.943969 1321573 gpu_helpers.cc:177] XLA backend will use up to 1275068416 bytes on device 0 for CollectiveBFCAllocator.
I0000 00:00:1745544722.958139 1321573 cuda_dnn.cc:527] Loaded cuDNN version 90400
E0000 00:00:1745545070.488441 1321573 buffer_comparator.cc:145] Difference at 16: 0, expected 11.6059
E0000 00:00:1745545070.488495 1321573 buffer_comparator.cc:145] Difference at 17: 0, expected 14.502
E0000 00:00:1745545070.488503 1321573 buffer_comparator.cc:145] Difference at 18: 0, expected 11.2449
E0000 00:00:1745545070.488509 1321573 buffer_comparator.cc:145] Difference at 19: 0, expected 10.0998
E0000 00:00:1745545070.488516 1321573 buffer_comparator.cc:145] Difference at 20: 0, expected 14.0222
E0000 00:00:1745545070.488523 1321573 buffer_comparator.cc:145] Difference at 21: 0, expected 10.1321
E0000 00:00:1745545070.488529 1321573 buffer_comparator.cc:145] Difference at 22: 0, expected 10.2986
E0000 00:00:1745545070.488536 1321573 buffer_comparator.cc:145] Difference at 23: 0, expected 14.1109
E0000 00:00:1745545070.488542 1321573 buffer_comparator.cc:145] Difference at 24: 0, expected 13.3463
E0000 00:00:1745545070.488549 1321573 buffer_comparator.cc:145] Difference at 25: 0, expected 12.8369
2025-04-25 01:37:50.488564: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.492138 1321573 buffer_comparator.cc:145] Difference at 16: 0, expected 11.6059
E0000 00:00:1745545070.492166 1321573 buffer_comparator.cc:145] Difference at 17: 0, expected 14.502
E0000 00:00:1745545070.492173 1321573 buffer_comparator.cc:145] Difference at 18: 0, expected 11.2449
E0000 00:00:1745545070.492180 1321573 buffer_comparator.cc:145] Difference at 19: 0, expected 10.0998
E0000 00:00:1745545070.492187 1321573 buffer_comparator.cc:145] Difference at 20: 0, expected 14.0222
E0000 00:00:1745545070.492193 1321573 buffer_comparator.cc:145] Difference at 21: 0, expected 10.1321
E0000 00:00:1745545070.492200 1321573 buffer_comparator.cc:145] Difference at 22: 0, expected 10.2986
E0000 00:00:1745545070.492206 1321573 buffer_comparator.cc:145] Difference at 23: 0, expected 14.1109
E0000 00:00:1745545070.492213 1321573 buffer_comparator.cc:145] Difference at 24: 0, expected 13.3463
E0000 00:00:1745545070.492219 1321573 buffer_comparator.cc:145] Difference at 25: 0, expected 12.8369
2025-04-25 01:37:50.492230: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.495654 1321573 buffer_comparator.cc:145] Difference at 16: 0, expected 11.6059
E0000 00:00:1745545070.495674 1321573 buffer_comparator.cc:145] Difference at 17: 0, expected 14.502
E0000 00:00:1745545070.495679 1321573 buffer_comparator.cc:145] Difference at 18: 0, expected 11.2449
E0000 00:00:1745545070.495683 1321573 buffer_comparator.cc:145] Difference at 19: 0, expected 10.0998
E0000 00:00:1745545070.495687 1321573 buffer_comparator.cc:145] Difference at 20: 0, expected 14.0222
E0000 00:00:1745545070.495692 1321573 buffer_comparator.cc:145] Difference at 21: 0, expected 10.1321
E0000 00:00:1745545070.495696 1321573 buffer_comparator.cc:145] Difference at 22: 0, expected 10.2986
E0000 00:00:1745545070.495701 1321573 buffer_comparator.cc:145] Difference at 23: 0, expected 14.1109
E0000 00:00:1745545070.495705 1321573 buffer_comparator.cc:145] Difference at 24: 0, expected 13.3463
E0000 00:00:1745545070.495711 1321573 buffer_comparator.cc:145] Difference at 25: 0, expected 12.8369
2025-04-25 01:37:50.495718: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.498805 1321573 buffer_comparator.cc:145] Difference at 32: 0, expected 12.4
E0000 00:00:1745545070.498827 1321573 buffer_comparator.cc:145] Difference at 33: 0, expected 12.9454
E0000 00:00:1745545070.498831 1321573 buffer_comparator.cc:145] Difference at 34: 0, expected 12.9462
E0000 00:00:1745545070.498836 1321573 buffer_comparator.cc:145] Difference at 35: 0, expected 13.9775
E0000 00:00:1745545070.498840 1321573 buffer_comparator.cc:145] Difference at 36: 0, expected 15.0433
E0000 00:00:1745545070.498844 1321573 buffer_comparator.cc:145] Difference at 37: 0, expected 12.0589
E0000 00:00:1745545070.498849 1321573 buffer_comparator.cc:145] Difference at 38: 0, expected 14.4629
E0000 00:00:1745545070.498853 1321573 buffer_comparator.cc:145] Difference at 39: 0, expected 12.7671
E0000 00:00:1745545070.498857 1321573 buffer_comparator.cc:145] Difference at 40: 0, expected 12.3584
E0000 00:00:1745545070.498862 1321573 buffer_comparator.cc:145] Difference at 41: 0, expected 11.6002
2025-04-25 01:37:50.498869: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.501946 1321573 buffer_comparator.cc:145] Difference at 1056: 0, expected 10.6543
E0000 00:00:1745545070.501968 1321573 buffer_comparator.cc:145] Difference at 1057: 0, expected 11.0945
E0000 00:00:1745545070.501973 1321573 buffer_comparator.cc:145] Difference at 1058: 0, expected 11.1424
E0000 00:00:1745545070.501977 1321573 buffer_comparator.cc:145] Difference at 1059: 0, expected 12.7556
E0000 00:00:1745545070.501982 1321573 buffer_comparator.cc:145] Difference at 1060: 0, expected 12.6932
E0000 00:00:1745545070.501986 1321573 buffer_comparator.cc:145] Difference at 1061: 0, expected 10.0594
E0000 00:00:1745545070.501990 1321573 buffer_comparator.cc:145] Difference at 1062: 0, expected 12.3478
E0000 00:00:1745545070.501995 1321573 buffer_comparator.cc:145] Difference at 1063: 0, expected 10.8381
E0000 00:00:1745545070.501999 1321573 buffer_comparator.cc:145] Difference at 1064: 0, expected 10.409
E0000 00:00:1745545070.502003 1321573 buffer_comparator.cc:145] Difference at 1065: 0, expected 10.3688
2025-04-25 01:37:50.502010: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.505085 1321573 buffer_comparator.cc:145] Difference at 1056: 0, expected 10.6543
E0000 00:00:1745545070.505104 1321573 buffer_comparator.cc:145] Difference at 1057: 0, expected 11.0945
E0000 00:00:1745545070.505109 1321573 buffer_comparator.cc:145] Difference at 1058: 0, expected 11.1424
E0000 00:00:1745545070.505113 1321573 buffer_comparator.cc:145] Difference at 1059: 0, expected 12.7556
E0000 00:00:1745545070.505117 1321573 buffer_comparator.cc:145] Difference at 1060: 0, expected 12.6932
E0000 00:00:1745545070.505122 1321573 buffer_comparator.cc:145] Difference at 1061: 0, expected 10.0594
E0000 00:00:1745545070.505126 1321573 buffer_comparator.cc:145] Difference at 1062: 0, expected 12.3478
E0000 00:00:1745545070.505130 1321573 buffer_comparator.cc:145] Difference at 1063: 0, expected 10.8381
E0000 00:00:1745545070.505135 1321573 buffer_comparator.cc:145] Difference at 1064: 0, expected 10.409
E0000 00:00:1745545070.505139 1321573 buffer_comparator.cc:145] Difference at 1065: 0, expected 10.3688
2025-04-25 01:37:50.505146: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.522754 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 9.17386
E0000 00:00:1745545070.522772 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 8.63962
E0000 00:00:1745545070.522775 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 11.0889
E0000 00:00:1745545070.522778 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 8.96333
E0000 00:00:1745545070.522781 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 8.01067
E0000 00:00:1745545070.522784 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 11.0277
E0000 00:00:1745545070.522786 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 7.61027
E0000 00:00:1745545070.522789 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 7.70495
E0000 00:00:1745545070.522792 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 10.9025
E0000 00:00:1745545070.522794 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 9.9377
2025-04-25 01:37:50.522799: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.524926 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 9.17386
E0000 00:00:1745545070.524942 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 8.63962
E0000 00:00:1745545070.524945 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 11.0889
E0000 00:00:1745545070.524948 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 8.96333
E0000 00:00:1745545070.524951 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 8.01067
E0000 00:00:1745545070.524953 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 11.0277
E0000 00:00:1745545070.524956 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 7.61027
E0000 00:00:1745545070.524959 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 7.70495
E0000 00:00:1745545070.524962 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 10.9025
E0000 00:00:1745545070.524964 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 9.9377
2025-04-25 01:37:50.524969: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.527110 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 9.17386
E0000 00:00:1745545070.527123 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 8.63962
E0000 00:00:1745545070.527126 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 11.0889
E0000 00:00:1745545070.527129 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 8.96333
E0000 00:00:1745545070.527131 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 8.01067
E0000 00:00:1745545070.527134 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 11.0277
E0000 00:00:1745545070.527137 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 7.61027
E0000 00:00:1745545070.527139 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 7.70495
E0000 00:00:1745545070.527142 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 10.9025
E0000 00:00:1745545070.527145 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 9.9377
2025-04-25 01:37:50.527149: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.529285 1321573 buffer_comparator.cc:145] Difference at 32: -nan, expected 9.26776
E0000 00:00:1745545070.529298 1321573 buffer_comparator.cc:145] Difference at 33: -nan, expected 9.1446
E0000 00:00:1745545070.529301 1321573 buffer_comparator.cc:145] Difference at 34: -nan, expected 9.79374
E0000 00:00:1745545070.529304 1321573 buffer_comparator.cc:145] Difference at 35: -nan, expected 10.23
E0000 00:00:1745545070.529308 1321573 buffer_comparator.cc:145] Difference at 36: -nan, expected 10.448
E0000 00:00:1745545070.529311 1321573 buffer_comparator.cc:145] Difference at 37: -nan, expected 11.8017
E0000 00:00:1745545070.529314 1321573 buffer_comparator.cc:145] Difference at 38: -nan, expected 9.69309
E0000 00:00:1745545070.529317 1321573 buffer_comparator.cc:145] Difference at 39: -nan, expected 11.2277
E0000 00:00:1745545070.529320 1321573 buffer_comparator.cc:145] Difference at 40: -nan, expected 10.0359
E0000 00:00:1745545070.529322 1321573 buffer_comparator.cc:145] Difference at 41: -nan, expected 9.30719
2025-04-25 01:37:50.529327: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.531466 1321573 buffer_comparator.cc:145] Difference at 1056: -nan, expected 6.66818
E0000 00:00:1745545070.531479 1321573 buffer_comparator.cc:145] Difference at 1057: -nan, expected 7.21702
E0000 00:00:1745545070.531482 1321573 buffer_comparator.cc:145] Difference at 1058: -nan, expected 7.51313
E0000 00:00:1745545070.531484 1321573 buffer_comparator.cc:145] Difference at 1059: -nan, expected 8.05613
E0000 00:00:1745545070.531487 1321573 buffer_comparator.cc:145] Difference at 1060: -nan, expected 9.84957
E0000 00:00:1745545070.531490 1321573 buffer_comparator.cc:145] Difference at 1061: -nan, expected 8.70676
E0000 00:00:1745545070.531493 1321573 buffer_comparator.cc:145] Difference at 1062: -nan, expected 6.83172
E0000 00:00:1745545070.531495 1321573 buffer_comparator.cc:145] Difference at 1063: -nan, expected 8.78604
E0000 00:00:1745545070.531498 1321573 buffer_comparator.cc:145] Difference at 1064: -nan, expected 7.62453
E0000 00:00:1745545070.531501 1321573 buffer_comparator.cc:145] Difference at 1065: -nan, expected 7.09446
2025-04-25 01:37:50.531505: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.533656 1321573 buffer_comparator.cc:145] Difference at 1056: -nan, expected 6.66818
E0000 00:00:1745545070.533668 1321573 buffer_comparator.cc:145] Difference at 1057: -nan, expected 7.21702
E0000 00:00:1745545070.533671 1321573 buffer_comparator.cc:145] Difference at 1058: -nan, expected 7.51313
E0000 00:00:1745545070.533674 1321573 buffer_comparator.cc:145] Difference at 1059: -nan, expected 8.05613
E0000 00:00:1745545070.533677 1321573 buffer_comparator.cc:145] Difference at 1060: -nan, expected 9.84957
E0000 00:00:1745545070.533679 1321573 buffer_comparator.cc:145] Difference at 1061: -nan, expected 8.70676
E0000 00:00:1745545070.533682 1321573 buffer_comparator.cc:145] Difference at 1062: -nan, expected 6.83172
E0000 00:00:1745545070.533685 1321573 buffer_comparator.cc:145] Difference at 1063: -nan, expected 8.78604
E0000 00:00:1745545070.533687 1321573 buffer_comparator.cc:145] Difference at 1064: -nan, expected 7.62453
E0000 00:00:1745545070.533690 1321573 buffer_comparator.cc:145] Difference at 1065: -nan, expected 7.09446
2025-04-25 01:37:50.533694: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.547116 1321573 buffer_comparator.cc:145] Difference at 16: 9.17386, expected 7.32342
E0000 00:00:1745545070.547129 1321573 buffer_comparator.cc:145] Difference at 17: 8.63962, expected 10.2674
E0000 00:00:1745545070.547132 1321573 buffer_comparator.cc:145] Difference at 18: 11.0889, expected 8.76472
E0000 00:00:1745545070.547136 1321573 buffer_comparator.cc:145] Difference at 21: 11.0277, expected 8.01524
E0000 00:00:1745545070.547139 1321573 buffer_comparator.cc:145] Difference at 22: 7.61027, expected 10.2221
E0000 00:00:1745545070.547142 1321573 buffer_comparator.cc:145] Difference at 24: 10.9025, expected 8.62185
E0000 00:00:1745545070.547145 1321573 buffer_comparator.cc:145] Difference at 29: 10.598, expected 8.06688
E0000 00:00:1745545070.547150 1321573 buffer_comparator.cc:145] Difference at 30: 10.6668, expected 9.17581
E0000 00:00:1745545070.547153 1321573 buffer_comparator.cc:145] Difference at 32: 9.26776, expected 7.79614
E0000 00:00:1745545070.547155 1321573 buffer_comparator.cc:145] Difference at 34: 9.79374, expected 7.73293
2025-04-25 01:37:50.547160: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.549295 1321573 buffer_comparator.cc:145] Difference at 16: 9.17386, expected 7.32342
E0000 00:00:1745545070.549309 1321573 buffer_comparator.cc:145] Difference at 17: 8.63962, expected 10.2674
E0000 00:00:1745545070.549312 1321573 buffer_comparator.cc:145] Difference at 18: 11.0889, expected 8.76472
E0000 00:00:1745545070.549315 1321573 buffer_comparator.cc:145] Difference at 21: 11.0277, expected 8.01524
E0000 00:00:1745545070.549318 1321573 buffer_comparator.cc:145] Difference at 22: 7.61027, expected 10.2221
E0000 00:00:1745545070.549321 1321573 buffer_comparator.cc:145] Difference at 24: 10.9025, expected 8.62185
E0000 00:00:1745545070.549324 1321573 buffer_comparator.cc:145] Difference at 29: 10.598, expected 8.06688
E0000 00:00:1745545070.549327 1321573 buffer_comparator.cc:145] Difference at 30: 10.6668, expected 9.17581
E0000 00:00:1745545070.549330 1321573 buffer_comparator.cc:145] Difference at 32: 9.26776, expected 7.79614
E0000 00:00:1745545070.549333 1321573 buffer_comparator.cc:145] Difference at 34: 9.79374, expected 7.73293
2025-04-25 01:37:50.549337: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.551472 1321573 buffer_comparator.cc:145] Difference at 16: 9.17386, expected 7.32342
E0000 00:00:1745545070.551490 1321573 buffer_comparator.cc:145] Difference at 17: 8.63962, expected 10.2674
E0000 00:00:1745545070.551494 1321573 buffer_comparator.cc:145] Difference at 18: 11.0889, expected 8.76472
E0000 00:00:1745545070.551497 1321573 buffer_comparator.cc:145] Difference at 21: 11.0277, expected 8.01524
E0000 00:00:1745545070.551500 1321573 buffer_comparator.cc:145] Difference at 22: 7.61027, expected 10.2221
E0000 00:00:1745545070.551503 1321573 buffer_comparator.cc:145] Difference at 24: 10.9025, expected 8.62185
E0000 00:00:1745545070.551506 1321573 buffer_comparator.cc:145] Difference at 29: 10.598, expected 8.06688
E0000 00:00:1745545070.551509 1321573 buffer_comparator.cc:145] Difference at 30: 10.6668, expected 9.17581
E0000 00:00:1745545070.551512 1321573 buffer_comparator.cc:145] Difference at 32: 9.26776, expected 7.79614
E0000 00:00:1745545070.551514 1321573 buffer_comparator.cc:145] Difference at 34: 9.79374, expected 7.73293
2025-04-25 01:37:50.551519: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.553660 1321573 buffer_comparator.cc:145] Difference at 32: 9.26776, expected 7.79614
E0000 00:00:1745545070.553673 1321573 buffer_comparator.cc:145] Difference at 34: 9.79374, expected 7.73293
E0000 00:00:1745545070.553676 1321573 buffer_comparator.cc:145] Difference at 36: 10.448, expected 8.80404
E0000 00:00:1745545070.553679 1321573 buffer_comparator.cc:145] Difference at 37: 11.8017, expected 10.1025
E0000 00:00:1745545070.553682 1321573 buffer_comparator.cc:145] Difference at 38: 9.69309, expected 8.25699
E0000 00:00:1745545070.553685 1321573 buffer_comparator.cc:145] Difference at 39: 11.2277, expected 8.00626
E0000 00:00:1745545070.553688 1321573 buffer_comparator.cc:145] Difference at 40: 10.0359, expected 8.86599
E0000 00:00:1745545070.553691 1321573 buffer_comparator.cc:145] Difference at 41: 9.30719, expected 7.81407
E0000 00:00:1745545070.553694 1321573 buffer_comparator.cc:145] Difference at 43: 10.7467, expected 8.32593
E0000 00:00:1745545070.553698 1321573 buffer_comparator.cc:145] Difference at 45: 9.80023, expected 8.35446
2025-04-25 01:37:50.553703: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.555846 1321573 buffer_comparator.cc:145] Difference at 1058: 7.51313, expected 6.62395
E0000 00:00:1745545070.555861 1321573 buffer_comparator.cc:145] Difference at 1060: 9.84957, expected 8.38545
E0000 00:00:1745545070.555864 1321573 buffer_comparator.cc:145] Difference at 1063: 8.78604, expected 7.14543
E0000 00:00:1745545070.555867 1321573 buffer_comparator.cc:145] Difference at 1065: 7.09446, expected 6.08633
E0000 00:00:1745545070.555869 1321573 buffer_comparator.cc:145] Difference at 1068: 5.91719, expected 8.19554
E0000 00:00:1745545070.555872 1321573 buffer_comparator.cc:145] Difference at 1071: 7.6916, expected 8.8658
E0000 00:00:1745545070.555875 1321573 buffer_comparator.cc:145] Difference at 1073: 8.42601, expected 6.83319
E0000 00:00:1745545070.555878 1321573 buffer_comparator.cc:145] Difference at 1074: 7.39256, expected 8.36791
E0000 00:00:1745545070.555881 1321573 buffer_comparator.cc:145] Difference at 1077: 6.50966, expected 8.08647
E0000 00:00:1745545070.555884 1321573 buffer_comparator.cc:145] Difference at 1078: 8.85918, expected 6.98496
2025-04-25 01:37:50.555888: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.558026 1321573 buffer_comparator.cc:145] Difference at 1058: 7.51313, expected 6.62395
E0000 00:00:1745545070.558038 1321573 buffer_comparator.cc:145] Difference at 1060: 9.84957, expected 8.38545
E0000 00:00:1745545070.558041 1321573 buffer_comparator.cc:145] Difference at 1063: 8.78604, expected 7.14543
E0000 00:00:1745545070.558044 1321573 buffer_comparator.cc:145] Difference at 1065: 7.09446, expected 6.08633
E0000 00:00:1745545070.558047 1321573 buffer_comparator.cc:145] Difference at 1068: 5.91719, expected 8.19554
E0000 00:00:1745545070.558057 1321573 buffer_comparator.cc:145] Difference at 1071: 7.6916, expected 8.8658
E0000 00:00:1745545070.558060 1321573 buffer_comparator.cc:145] Difference at 1073: 8.42601, expected 6.83319
E0000 00:00:1745545070.558062 1321573 buffer_comparator.cc:145] Difference at 1074: 7.39256, expected 8.36791
E0000 00:00:1745545070.558065 1321573 buffer_comparator.cc:145] Difference at 1077: 6.50966, expected 8.08647
E0000 00:00:1745545070.558068 1321573 buffer_comparator.cc:145] Difference at 1078: 8.85918, expected 6.98496
2025-04-25 01:37:50.558073: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.586422 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 7.73985
E0000 00:00:1745545070.586456 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 7.11199
E0000 00:00:1745545070.586459 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 7.78251
E0000 00:00:1745545070.586461 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 8.01723
E0000 00:00:1745545070.586464 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 7.54006
E0000 00:00:1745545070.586467 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 9.39119
E0000 00:00:1745545070.586470 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 9.38334
E0000 00:00:1745545070.586472 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 7.26531
E0000 00:00:1745545070.586475 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 8.44799
E0000 00:00:1745545070.586478 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 7.82992
2025-04-25 01:37:50.586485: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.588625 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 7.73985
E0000 00:00:1745545070.588636 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 7.11199
E0000 00:00:1745545070.588639 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 7.78251
E0000 00:00:1745545070.588642 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 8.01723
E0000 00:00:1745545070.588644 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 7.54006
E0000 00:00:1745545070.588647 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 9.39119
E0000 00:00:1745545070.588650 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 9.38334
E0000 00:00:1745545070.588652 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 7.26531
E0000 00:00:1745545070.588655 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 8.44799
E0000 00:00:1745545070.588658 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 7.82992
2025-04-25 01:37:50.588662: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.590797 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 7.73985
E0000 00:00:1745545070.590811 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 7.11199
E0000 00:00:1745545070.590814 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 7.78251
E0000 00:00:1745545070.590816 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 8.01723
E0000 00:00:1745545070.590819 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 7.54006
E0000 00:00:1745545070.590822 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 9.39119
E0000 00:00:1745545070.590825 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 9.38334
E0000 00:00:1745545070.590827 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 7.26531
E0000 00:00:1745545070.590830 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 8.44799
E0000 00:00:1745545070.590833 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 7.82992
2025-04-25 01:37:50.590837: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.592969 1321573 buffer_comparator.cc:145] Difference at 32: -nan, expected 9.09096
E0000 00:00:1745545070.592985 1321573 buffer_comparator.cc:145] Difference at 33: -nan, expected 9.0377
E0000 00:00:1745545070.592988 1321573 buffer_comparator.cc:145] Difference at 34: -nan, expected 7.48572
E0000 00:00:1745545070.592990 1321573 buffer_comparator.cc:145] Difference at 35: -nan, expected 9.05382
E0000 00:00:1745545070.592993 1321573 buffer_comparator.cc:145] Difference at 36: -nan, expected 8.75347
E0000 00:00:1745545070.592996 1321573 buffer_comparator.cc:145] Difference at 37: -nan, expected 7.34985
E0000 00:00:1745545070.592998 1321573 buffer_comparator.cc:145] Difference at 38: -nan, expected 8.78563
E0000 00:00:1745545070.593001 1321573 buffer_comparator.cc:145] Difference at 39: -nan, expected 7.6147
E0000 00:00:1745545070.593004 1321573 buffer_comparator.cc:145] Difference at 40: -nan, expected 6.04756
E0000 00:00:1745545070.593006 1321573 buffer_comparator.cc:145] Difference at 41: -nan, expected 7.8893
2025-04-25 01:37:50.593011: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.595150 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.43679
E0000 00:00:1745545070.595167 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 9.73984
E0000 00:00:1745545070.595171 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.0838
E0000 00:00:1745545070.595175 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.70695
E0000 00:00:1745545070.595177 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.00042
E0000 00:00:1745545070.595180 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 8.345
E0000 00:00:1745545070.595183 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 8.52479
E0000 00:00:1745545070.595185 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 7.08541
E0000 00:00:1745545070.595188 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 8.47292
E0000 00:00:1745545070.595190 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 8.09814
2025-04-25 01:37:50.595195: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.597320 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.43679
E0000 00:00:1745545070.597334 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 9.73984
E0000 00:00:1745545070.597337 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.0838
E0000 00:00:1745545070.597340 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.70695
E0000 00:00:1745545070.597342 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.00042
E0000 00:00:1745545070.597345 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 8.345
E0000 00:00:1745545070.597348 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 8.52479
E0000 00:00:1745545070.597350 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 7.08541
E0000 00:00:1745545070.597353 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 8.47292
E0000 00:00:1745545070.597356 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 8.09814
2025-04-25 01:37:50.597360: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.599492 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.43679
E0000 00:00:1745545070.599506 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 9.73984
E0000 00:00:1745545070.599509 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.0838
E0000 00:00:1745545070.599512 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.70695
E0000 00:00:1745545070.599515 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.00042
E0000 00:00:1745545070.599517 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 8.345
E0000 00:00:1745545070.599520 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 8.52479
E0000 00:00:1745545070.599523 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 7.08541
E0000 00:00:1745545070.599525 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 8.47292
E0000 00:00:1745545070.599528 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 8.09814
2025-04-25 01:37:50.599533: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.601654 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.43679
E0000 00:00:1745545070.601670 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 9.73984
E0000 00:00:1745545070.601673 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.0838
E0000 00:00:1745545070.601676 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.70695
E0000 00:00:1745545070.601679 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.00042
E0000 00:00:1745545070.601683 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 8.345
E0000 00:00:1745545070.601686 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 8.52479
E0000 00:00:1745545070.601688 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 7.08541
E0000 00:00:1745545070.601691 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 8.47292
E0000 00:00:1745545070.601694 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 8.09814
2025-04-25 01:37:50.601698: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.603837 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.43679
E0000 00:00:1745545070.603852 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 9.73984
E0000 00:00:1745545070.603855 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.0838
E0000 00:00:1745545070.603858 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.70695
E0000 00:00:1745545070.603860 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.00042
E0000 00:00:1745545070.603863 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 8.345
E0000 00:00:1745545070.603866 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 8.52479
E0000 00:00:1745545070.603868 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 7.08541
E0000 00:00:1745545070.603871 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 8.47292
E0000 00:00:1745545070.603874 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 8.09814
2025-04-25 01:37:50.603878: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.605649 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.43679
E0000 00:00:1745545070.605663 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 9.73984
E0000 00:00:1745545070.605666 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.0838
E0000 00:00:1745545070.605669 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.70695
E0000 00:00:1745545070.605671 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.00042
E0000 00:00:1745545070.605674 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 8.345
E0000 00:00:1745545070.605677 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 8.52479
E0000 00:00:1745545070.605680 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 7.08541
E0000 00:00:1745545070.605682 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 8.47292
E0000 00:00:1745545070.605685 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 8.09814
2025-04-25 01:37:50.605689: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.607423 1321573 buffer_comparator.cc:145] Difference at 128: -nan, expected 10.3054
E0000 00:00:1745545070.607438 1321573 buffer_comparator.cc:145] Difference at 129: -nan, expected 9.87733
E0000 00:00:1745545070.607441 1321573 buffer_comparator.cc:145] Difference at 130: -nan, expected 8.66034
E0000 00:00:1745545070.607444 1321573 buffer_comparator.cc:145] Difference at 131: -nan, expected 8.49512
E0000 00:00:1745545070.607446 1321573 buffer_comparator.cc:145] Difference at 132: -nan, expected 9.08348
E0000 00:00:1745545070.607449 1321573 buffer_comparator.cc:145] Difference at 133: -nan, expected 9.79527
E0000 00:00:1745545070.607452 1321573 buffer_comparator.cc:145] Difference at 134: -nan, expected 7.65675
E0000 00:00:1745545070.607454 1321573 buffer_comparator.cc:145] Difference at 135: -nan, expected 7.69855
E0000 00:00:1745545070.607459 1321573 buffer_comparator.cc:145] Difference at 136: -nan, expected 8.6248
E0000 00:00:1745545070.607461 1321573 buffer_comparator.cc:145] Difference at 137: -nan, expected 8.67269
2025-04-25 01:37:50.607466: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.609816 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 8.81126
E0000 00:00:1745545070.609827 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 10.499
E0000 00:00:1745545070.609830 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 9.85809
E0000 00:00:1745545070.609833 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 8.65377
E0000 00:00:1745545070.609836 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 10.0279
E0000 00:00:1745545070.609839 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 9.616
E0000 00:00:1745545070.609841 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 11.4222
E0000 00:00:1745545070.609844 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 8.99654
E0000 00:00:1745545070.609847 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 10.1518
E0000 00:00:1745545070.609849 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 9.78295
2025-04-25 01:37:50.609854: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.611580 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 8.81126
E0000 00:00:1745545070.611594 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 10.499
E0000 00:00:1745545070.611597 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 9.85809
E0000 00:00:1745545070.611600 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 8.65377
E0000 00:00:1745545070.611603 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 10.0279
E0000 00:00:1745545070.611605 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 9.616
E0000 00:00:1745545070.611608 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 11.4222
E0000 00:00:1745545070.611611 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 8.99654
E0000 00:00:1745545070.611614 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 10.1518
E0000 00:00:1745545070.611616 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 9.78295
2025-04-25 01:37:50.611621: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.613366 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 8.81126
E0000 00:00:1745545070.613381 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 10.499
E0000 00:00:1745545070.613384 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 9.85809
E0000 00:00:1745545070.613387 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 8.65377
E0000 00:00:1745545070.613390 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 10.0279
E0000 00:00:1745545070.613392 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 9.616
E0000 00:00:1745545070.613395 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 11.4222
E0000 00:00:1745545070.613398 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 8.99654
E0000 00:00:1745545070.613401 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 10.1518
E0000 00:00:1745545070.613403 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 9.78295
2025-04-25 01:37:50.613408: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.615153 1321573 buffer_comparator.cc:145] Difference at 32: -nan, expected 8.20589
E0000 00:00:1745545070.615169 1321573 buffer_comparator.cc:145] Difference at 33: -nan, expected 9.19448
E0000 00:00:1745545070.615172 1321573 buffer_comparator.cc:145] Difference at 34: -nan, expected 8.03576
E0000 00:00:1745545070.615174 1321573 buffer_comparator.cc:145] Difference at 35: -nan, expected 10.6612
E0000 00:00:1745545070.615177 1321573 buffer_comparator.cc:145] Difference at 36: -nan, expected 9.50115
E0000 00:00:1745545070.615180 1321573 buffer_comparator.cc:145] Difference at 37: -nan, expected 10.8965
E0000 00:00:1745545070.615183 1321573 buffer_comparator.cc:145] Difference at 38: -nan, expected 9.00596
E0000 00:00:1745545070.615186 1321573 buffer_comparator.cc:145] Difference at 39: -nan, expected 7.25269
E0000 00:00:1745545070.615188 1321573 buffer_comparator.cc:145] Difference at 40: -nan, expected 9.67287
E0000 00:00:1745545070.615191 1321573 buffer_comparator.cc:145] Difference at 41: -nan, expected 7.30318
2025-04-25 01:37:50.615195: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.616941 1321573 buffer_comparator.cc:145] Difference at 1056: -nan, expected 6.05049
E0000 00:00:1745545070.616955 1321573 buffer_comparator.cc:145] Difference at 1057: -nan, expected 7.30224
E0000 00:00:1745545070.616958 1321573 buffer_comparator.cc:145] Difference at 1058: -nan, expected 7.29652
E0000 00:00:1745545070.616961 1321573 buffer_comparator.cc:145] Difference at 1059: -nan, expected 8.69607
E0000 00:00:1745545070.616964 1321573 buffer_comparator.cc:145] Difference at 1060: -nan, expected 7.43241
E0000 00:00:1745545070.616966 1321573 buffer_comparator.cc:145] Difference at 1061: -nan, expected 7.71511
E0000 00:00:1745545070.616969 1321573 buffer_comparator.cc:145] Difference at 1062: -nan, expected 7.76911
E0000 00:00:1745545070.616972 1321573 buffer_comparator.cc:145] Difference at 1063: -nan, expected 7.29298
E0000 00:00:1745545070.616974 1321573 buffer_comparator.cc:145] Difference at 1064: -nan, expected 7.0298
E0000 00:00:1745545070.616977 1321573 buffer_comparator.cc:145] Difference at 1065: -nan, expected 5.85707
2025-04-25 01:37:50.616981: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.618723 1321573 buffer_comparator.cc:145] Difference at 1056: -nan, expected 6.05049
E0000 00:00:1745545070.618736 1321573 buffer_comparator.cc:145] Difference at 1057: -nan, expected 7.30224
E0000 00:00:1745545070.618739 1321573 buffer_comparator.cc:145] Difference at 1058: -nan, expected 7.29652
E0000 00:00:1745545070.618742 1321573 buffer_comparator.cc:145] Difference at 1059: -nan, expected 8.69607
E0000 00:00:1745545070.618744 1321573 buffer_comparator.cc:145] Difference at 1060: -nan, expected 7.43241
E0000 00:00:1745545070.618747 1321573 buffer_comparator.cc:145] Difference at 1061: -nan, expected 7.71511
E0000 00:00:1745545070.618750 1321573 buffer_comparator.cc:145] Difference at 1062: -nan, expected 7.76911
E0000 00:00:1745545070.618752 1321573 buffer_comparator.cc:145] Difference at 1063: -nan, expected 7.29298
E0000 00:00:1745545070.618755 1321573 buffer_comparator.cc:145] Difference at 1064: -nan, expected 7.0298
E0000 00:00:1745545070.618758 1321573 buffer_comparator.cc:145] Difference at 1065: -nan, expected 5.85707
2025-04-25 01:37:50.618762: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.632125 1321573 buffer_comparator.cc:145] Difference at 1024: -nan, expected 0.527629
E0000 00:00:1745545070.632155 1321573 buffer_comparator.cc:145] Difference at 1025: -nan, expected 0.544573
E0000 00:00:1745545070.632160 1321573 buffer_comparator.cc:145] Difference at 1026: -nan, expected 0.31953
E0000 00:00:1745545070.632162 1321573 buffer_comparator.cc:145] Difference at 1027: -nan, expected 0.61527
E0000 00:00:1745545070.632165 1321573 buffer_comparator.cc:145] Difference at 1028: -nan, expected 1.06241
E0000 00:00:1745545070.632168 1321573 buffer_comparator.cc:145] Difference at 1029: -nan, expected 1.22093
E0000 00:00:1745545070.632171 1321573 buffer_comparator.cc:145] Difference at 1030: -nan, expected 0.553754
E0000 00:00:1745545070.632173 1321573 buffer_comparator.cc:145] Difference at 1031: -nan, expected 0.665092
E0000 00:00:1745545070.632176 1321573 buffer_comparator.cc:145] Difference at 1032: -nan, expected 1.13956
E0000 00:00:1745545070.632179 1321573 buffer_comparator.cc:145] Difference at 1033: -nan, expected 0.481583
2025-04-25 01:37:50.632186: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.643873 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 9.278
E0000 00:00:1745545070.643914 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 9.69724
E0000 00:00:1745545070.643917 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 10.7774
E0000 00:00:1745545070.643920 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 10.3636
E0000 00:00:1745545070.643923 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 7.56533
E0000 00:00:1745545070.643926 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 8.17756
E0000 00:00:1745545070.643928 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 11.0551
E0000 00:00:1745545070.643931 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 11.628
E0000 00:00:1745545070.643934 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 10.6715
E0000 00:00:1745545070.643937 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 10.8131
2025-04-25 01:37:50.643946: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.645700 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 9.278
E0000 00:00:1745545070.645716 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 9.69724
E0000 00:00:1745545070.645719 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 10.7774
E0000 00:00:1745545070.645722 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 10.3636
E0000 00:00:1745545070.645725 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 7.56533
E0000 00:00:1745545070.645728 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 8.17756
E0000 00:00:1745545070.645730 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 11.0551
E0000 00:00:1745545070.645733 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 11.628
E0000 00:00:1745545070.645736 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 10.6715
E0000 00:00:1745545070.645739 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 10.8131
2025-04-25 01:37:50.645743: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.647499 1321573 buffer_comparator.cc:145] Difference at 16: -nan, expected 9.278
E0000 00:00:1745545070.647513 1321573 buffer_comparator.cc:145] Difference at 17: -nan, expected 9.69724
E0000 00:00:1745545070.647515 1321573 buffer_comparator.cc:145] Difference at 18: -nan, expected 10.7774
E0000 00:00:1745545070.647519 1321573 buffer_comparator.cc:145] Difference at 19: -nan, expected 10.3636
E0000 00:00:1745545070.647521 1321573 buffer_comparator.cc:145] Difference at 20: -nan, expected 7.56533
E0000 00:00:1745545070.647526 1321573 buffer_comparator.cc:145] Difference at 21: -nan, expected 8.17756
E0000 00:00:1745545070.647528 1321573 buffer_comparator.cc:145] Difference at 22: -nan, expected 11.0551
E0000 00:00:1745545070.647531 1321573 buffer_comparator.cc:145] Difference at 23: -nan, expected 11.628
E0000 00:00:1745545070.647534 1321573 buffer_comparator.cc:145] Difference at 24: -nan, expected 10.6715
E0000 00:00:1745545070.647537 1321573 buffer_comparator.cc:145] Difference at 25: -nan, expected 10.8131
2025-04-25 01:37:50.647541: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.649290 1321573 buffer_comparator.cc:145] Difference at 32: -nan, expected 10.5897
E0000 00:00:1745545070.649303 1321573 buffer_comparator.cc:145] Difference at 33: -nan, expected 8.79857
E0000 00:00:1745545070.649306 1321573 buffer_comparator.cc:145] Difference at 34: -nan, expected 10.7009
E0000 00:00:1745545070.649309 1321573 buffer_comparator.cc:145] Difference at 35: -nan, expected 10.0423
E0000 00:00:1745545070.649312 1321573 buffer_comparator.cc:145] Difference at 36: -nan, expected 10.3055
E0000 00:00:1745545070.649315 1321573 buffer_comparator.cc:145] Difference at 37: -nan, expected 11.1027
E0000 00:00:1745545070.649318 1321573 buffer_comparator.cc:145] Difference at 38: -nan, expected 9.64939
E0000 00:00:1745545070.649320 1321573 buffer_comparator.cc:145] Difference at 39: -nan, expected 10.2558
E0000 00:00:1745545070.649323 1321573 buffer_comparator.cc:145] Difference at 40: -nan, expected 8.38055
E0000 00:00:1745545070.649326 1321573 buffer_comparator.cc:145] Difference at 41: -nan, expected 7.92221
2025-04-25 01:37:50.649330: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.651072 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.4725
E0000 00:00:1745545070.651084 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 8.00729
E0000 00:00:1745545070.651087 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.5558
E0000 00:00:1745545070.651090 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.97013
E0000 00:00:1745545070.651093 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.17677
E0000 00:00:1745545070.651096 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 6.91094
E0000 00:00:1745545070.651098 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 7.85311
E0000 00:00:1745545070.651101 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 8.69839
E0000 00:00:1745545070.651104 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 9.37233
E0000 00:00:1745545070.651106 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 7.56939
2025-04-25 01:37:50.651111: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.652863 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.4725
E0000 00:00:1745545070.652876 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 8.00729
E0000 00:00:1745545070.652878 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.5558
E0000 00:00:1745545070.652881 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.97013
E0000 00:00:1745545070.652884 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.17677
E0000 00:00:1745545070.652887 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 6.91094
E0000 00:00:1745545070.652889 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 7.85311
E0000 00:00:1745545070.652892 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 8.69839
E0000 00:00:1745545070.652896 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 9.37233
E0000 00:00:1745545070.652899 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 7.56939
2025-04-25 01:37:50.652904: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.654687 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.4725
E0000 00:00:1745545070.654701 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 8.00729
E0000 00:00:1745545070.654704 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.5558
E0000 00:00:1745545070.654707 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.97013
E0000 00:00:1745545070.654709 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.17677
E0000 00:00:1745545070.654712 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 6.91094
E0000 00:00:1745545070.654715 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 7.85311
E0000 00:00:1745545070.654717 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 8.69839
E0000 00:00:1745545070.654720 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 9.37233
E0000 00:00:1745545070.654723 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 7.56939
2025-04-25 01:37:50.654727: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.656475 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.4725
E0000 00:00:1745545070.656490 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 8.00729
E0000 00:00:1745545070.656493 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.5558
E0000 00:00:1745545070.656496 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.97013
E0000 00:00:1745545070.656499 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.17677
E0000 00:00:1745545070.656502 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 6.91094
E0000 00:00:1745545070.656504 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 7.85311
E0000 00:00:1745545070.656507 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 8.69839
E0000 00:00:1745545070.656509 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 9.37233
E0000 00:00:1745545070.656512 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 7.56939
2025-04-25 01:37:50.656517: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.658269 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.4725
E0000 00:00:1745545070.658282 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 8.00729
E0000 00:00:1745545070.658285 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.5558
E0000 00:00:1745545070.658288 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.97013
E0000 00:00:1745545070.658291 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.17677
E0000 00:00:1745545070.658294 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 6.91094
E0000 00:00:1745545070.658296 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 7.85311
E0000 00:00:1745545070.658299 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 8.69839
E0000 00:00:1745545070.658302 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 9.37233
E0000 00:00:1745545070.658304 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 7.56939
2025-04-25 01:37:50.658309: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0000 00:00:1745545070.660053 1321573 buffer_comparator.cc:145] Difference at 64: -nan, expected 8.4725
E0000 00:00:1745545070.660068 1321573 buffer_comparator.cc:145] Difference at 65: -nan, expected 8.00729
E0000 00:00:1745545070.660071 1321573 buffer_comparator.cc:145] Difference at 66: -nan, expected 10.5558
E0000 00:00:1745545070.660075 1321573 buffer_comparator.cc:145] Difference at 67: -nan, expected 7.97013
E0000 00:00:1745545070.660077 1321573 buffer_comparator.cc:145] Difference at 68: -nan, expected 9.17677
E0000 00:00:1745545070.660080 1321573 buffer_comparator.cc:145] Difference at 69: -nan, expected 6.91094
E0000 00:00:1745545070.660082 1321573 buffer_comparator.cc:145] Difference at 70: -nan, expected 7.85311
E0000 00:00:1745545070.660085 1321573 buffer_comparator.cc:145] Difference at 71: -nan, expected 8.69839
E0000 00:00:1745545070.660088 1321573 buffer_comparator.cc:145] Difference at 72: -nan, expected 9.37233
E0000 00:00:1745545070.660091 1321573 buffer_comparator.cc:145] Difference at 73: -nan, expected 7.56939
2025-04-25 01:37:50.660095: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1172] 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.034476429 (0.028613850) 	 Physics Loss: 0.000537277 (0.000464591) 	 Data Loss: 0.022220680 (0.011835133) 	 BC Loss: 0.011718473 (0.016314128)
Iteration: [  2001/ 50000] 	 Loss: 0.028288107 (0.030332621) 	 Physics Loss: 0.001888818 (0.001332523) 	 Data Loss: 0.012324474 (0.011986902) 	 BC Loss: 0.014074816 (0.017013200)
Iteration: [  3001/ 50000] 	 Loss: 0.022475120 (0.029750597) 	 Physics Loss: 0.000509278 (0.005786100) 	 Data Loss: 0.004655085 (0.008586247) 	 BC Loss: 0.017310757 (0.015378252)
Iteration: [  4001/ 50000] 	 Loss: 0.051963162 (0.038690042) 	 Physics Loss: 0.000808854 (0.002378667) 	 Data Loss: 0.020114277 (0.014395618) 	 BC Loss: 0.031040031 (0.021915752)
Iteration: [  5001/ 50000] 	 Loss: 0.022607081 (0.030625438) 	 Physics Loss: 0.007139057 (0.004731141) 	 Data Loss: 0.007296032 (0.010895450) 	 BC Loss: 0.008171991 (0.014998848)
Iteration: [  6001/ 50000] 	 Loss: 0.023505628 (0.019416459) 	 Physics Loss: 0.000467509 (0.000675986) 	 Data Loss: 0.005958247 (0.006814920) 	 BC Loss: 0.017079871 (0.011925554)
Iteration: [  7001/ 50000] 	 Loss: 0.017825369 (0.020186940) 	 Physics Loss: 0.000997673 (0.000751859) 	 Data Loss: 0.006823230 (0.008011221) 	 BC Loss: 0.010004465 (0.011423863)
Iteration: [  8001/ 50000] 	 Loss: 0.018535037 (0.017670767) 	 Physics Loss: 0.000696681 (0.001136744) 	 Data Loss: 0.004018028 (0.005713964) 	 BC Loss: 0.013820328 (0.010820058)
Iteration: [  9001/ 50000] 	 Loss: 0.019631973 (0.016841047) 	 Physics Loss: 0.000899504 (0.001499510) 	 Data Loss: 0.003037816 (0.005122127) 	 BC Loss: 0.015694652 (0.010219410)
Iteration: [ 10001/ 50000] 	 Loss: 0.012481965 (0.016115462) 	 Physics Loss: 0.001286428 (0.001695180) 	 Data Loss: 0.005454580 (0.004960962) 	 BC Loss: 0.005740956 (0.009459320)
Iteration: [ 11001/ 50000] 	 Loss: 0.017406404 (0.014580002) 	 Physics Loss: 0.001195281 (0.001443345) 	 Data Loss: 0.004605423 (0.004137451) 	 BC Loss: 0.011605700 (0.008999207)
Iteration: [ 12001/ 50000] 	 Loss: 0.014652768 (0.013691602) 	 Physics Loss: 0.002218292 (0.001398625) 	 Data Loss: 0.006017622 (0.003945971) 	 BC Loss: 0.006416854 (0.008347007)
Iteration: [ 13001/ 50000] 	 Loss: 0.007062657 (0.012623465) 	 Physics Loss: 0.001060225 (0.001372684) 	 Data Loss: 0.001934168 (0.003288924) 	 BC Loss: 0.004068264 (0.007961855)
Iteration: [ 14001/ 50000] 	 Loss: 0.007603331 (0.010438737) 	 Physics Loss: 0.001233277 (0.001874057) 	 Data Loss: 0.003455219 (0.002572269) 	 BC Loss: 0.002914835 (0.005992410)
Iteration: [ 15001/ 50000] 	 Loss: 0.008004644 (0.008404493) 	 Physics Loss: 0.001387044 (0.002349057) 	 Data Loss: 0.000460023 (0.002321562) 	 BC Loss: 0.006157577 (0.003733875)
Iteration: [ 16001/ 50000] 	 Loss: 0.005613001 (0.007259070) 	 Physics Loss: 0.001843056 (0.002004278) 	 Data Loss: 0.002583707 (0.001728758) 	 BC Loss: 0.001186238 (0.003526034)
Iteration: [ 17001/ 50000] 	 Loss: 0.003647973 (0.006060385) 	 Physics Loss: 0.001992107 (0.001862012) 	 Data Loss: 0.000419229 (0.001467435) 	 BC Loss: 0.001236636 (0.002730938)
Iteration: [ 18001/ 50000] 	 Loss: 0.006655810 (0.005444570) 	 Physics Loss: 0.002263136 (0.001785936) 	 Data Loss: 0.001224814 (0.001281822) 	 BC Loss: 0.003167860 (0.002376811)
Iteration: [ 19001/ 50000] 	 Loss: 0.004345927 (0.005763490) 	 Physics Loss: 0.001127900 (0.001784010) 	 Data Loss: 0.001592375 (0.001493493) 	 BC Loss: 0.001625652 (0.002485987)
Iteration: [ 20001/ 50000] 	 Loss: 0.005687185 (0.005456816) 	 Physics Loss: 0.002990145 (0.001652343) 	 Data Loss: 0.002053361 (0.001223832) 	 BC Loss: 0.000643680 (0.002580642)
Iteration: [ 21001/ 50000] 	 Loss: 0.004015730 (0.005559377) 	 Physics Loss: 0.001204707 (0.001961588) 	 Data Loss: 0.001680051 (0.001537366) 	 BC Loss: 0.001130971 (0.002060424)
Iteration: [ 22001/ 50000] 	 Loss: 0.006090690 (0.005603793) 	 Physics Loss: 0.004578053 (0.001906999) 	 Data Loss: 0.001199006 (0.001238090) 	 BC Loss: 0.000313632 (0.002458705)
Iteration: [ 23001/ 50000] 	 Loss: 0.007875150 (0.004338282) 	 Physics Loss: 0.002729153 (0.001461424) 	 Data Loss: 0.001961816 (0.001220324) 	 BC Loss: 0.003184182 (0.001656534)
Iteration: [ 24001/ 50000] 	 Loss: 0.004983329 (0.004078004) 	 Physics Loss: 0.001170014 (0.001519648) 	 Data Loss: 0.000645940 (0.000871195) 	 BC Loss: 0.003167375 (0.001687160)
Iteration: [ 25001/ 50000] 	 Loss: 0.002843466 (0.003566506) 	 Physics Loss: 0.001135944 (0.001550639) 	 Data Loss: 0.000395693 (0.000841146) 	 BC Loss: 0.001311828 (0.001174721)
Iteration: [ 26001/ 50000] 	 Loss: 0.003841943 (0.003158615) 	 Physics Loss: 0.001410076 (0.001508491) 	 Data Loss: 0.000356743 (0.000717544) 	 BC Loss: 0.002075125 (0.000932580)
Iteration: [ 27001/ 50000] 	 Loss: 0.003060500 (0.002717576) 	 Physics Loss: 0.000701304 (0.001169382) 	 Data Loss: 0.000589840 (0.000671310) 	 BC Loss: 0.001769357 (0.000876884)
Iteration: [ 28001/ 50000] 	 Loss: 0.002082543 (0.003386450) 	 Physics Loss: 0.001111371 (0.002014488) 	 Data Loss: 0.000513317 (0.000673456) 	 BC Loss: 0.000457855 (0.000698507)
Iteration: [ 29001/ 50000] 	 Loss: 0.001967198 (0.002122816) 	 Physics Loss: 0.000973968 (0.000992660) 	 Data Loss: 0.000552619 (0.000568144) 	 BC Loss: 0.000440611 (0.000562012)
Iteration: [ 30001/ 50000] 	 Loss: 0.002126399 (0.002069690) 	 Physics Loss: 0.000664904 (0.000953037) 	 Data Loss: 0.000561481 (0.000576463) 	 BC Loss: 0.000900015 (0.000540189)
Iteration: [ 31001/ 50000] 	 Loss: 0.002431013 (0.002888523) 	 Physics Loss: 0.000977250 (0.001792691) 	 Data Loss: 0.000422885 (0.000558199) 	 BC Loss: 0.001030877 (0.000537633)
Iteration: [ 32001/ 50000] 	 Loss: 0.002969190 (0.002110820) 	 Physics Loss: 0.002168338 (0.001169205) 	 Data Loss: 0.000308801 (0.000537486) 	 BC Loss: 0.000492052 (0.000404129)
Iteration: [ 33001/ 50000] 	 Loss: 0.001410649 (0.001557581) 	 Physics Loss: 0.000613028 (0.000788432) 	 Data Loss: 0.000337314 (0.000426982) 	 BC Loss: 0.000460307 (0.000342167)
Iteration: [ 34001/ 50000] 	 Loss: 0.001038901 (0.001468005) 	 Physics Loss: 0.000461820 (0.000700260) 	 Data Loss: 0.000267806 (0.000456997) 	 BC Loss: 0.000309275 (0.000310748)
Iteration: [ 35001/ 50000] 	 Loss: 0.000676396 (0.001624404) 	 Physics Loss: 0.000394395 (0.000935485) 	 Data Loss: 0.000131226 (0.000407618) 	 BC Loss: 0.000150774 (0.000281301)
Iteration: [ 36001/ 50000] 	 Loss: 0.002317710 (0.001326067) 	 Physics Loss: 0.001039488 (0.000754847) 	 Data Loss: 0.001045658 (0.000358247) 	 BC Loss: 0.000232563 (0.000212974)
Iteration: [ 37001/ 50000] 	 Loss: 0.001742344 (0.001391601) 	 Physics Loss: 0.001244900 (0.000813981) 	 Data Loss: 0.000260421 (0.000377276) 	 BC Loss: 0.000237024 (0.000200344)
Iteration: [ 38001/ 50000] 	 Loss: 0.001430800 (0.001394306) 	 Physics Loss: 0.000597752 (0.000799754) 	 Data Loss: 0.000443046 (0.000381376) 	 BC Loss: 0.000390002 (0.000213176)
Iteration: [ 39001/ 50000] 	 Loss: 0.001238087 (0.001228895) 	 Physics Loss: 0.000659979 (0.000639257) 	 Data Loss: 0.000283394 (0.000403422) 	 BC Loss: 0.000294714 (0.000186216)
Iteration: [ 40001/ 50000] 	 Loss: 0.002127211 (0.001517499) 	 Physics Loss: 0.001483771 (0.000949340) 	 Data Loss: 0.000404907 (0.000377177) 	 BC Loss: 0.000238533 (0.000190982)
Iteration: [ 41001/ 50000] 	 Loss: 0.000941515 (0.001476692) 	 Physics Loss: 0.000643017 (0.000909959) 	 Data Loss: 0.000242894 (0.000402896) 	 BC Loss: 0.000055604 (0.000163837)
Iteration: [ 42001/ 50000] 	 Loss: 0.002060887 (0.001324099) 	 Physics Loss: 0.001905333 (0.000875430) 	 Data Loss: 0.000100745 (0.000310349) 	 BC Loss: 0.000054809 (0.000138320)
Iteration: [ 43001/ 50000] 	 Loss: 0.000986727 (0.001232710) 	 Physics Loss: 0.000619240 (0.000712482) 	 Data Loss: 0.000276834 (0.000321609) 	 BC Loss: 0.000090654 (0.000198618)
Iteration: [ 44001/ 50000] 	 Loss: 0.000783800 (0.001046131) 	 Physics Loss: 0.000399274 (0.000563082) 	 Data Loss: 0.000231505 (0.000337236) 	 BC Loss: 0.000153021 (0.000145813)
Iteration: [ 45001/ 50000] 	 Loss: 0.000944633 (0.001041434) 	 Physics Loss: 0.000353833 (0.000624185) 	 Data Loss: 0.000500507 (0.000299698) 	 BC Loss: 0.000090293 (0.000117550)
Iteration: [ 46001/ 50000] 	 Loss: 0.001212354 (0.001046339) 	 Physics Loss: 0.000637281 (0.000556225) 	 Data Loss: 0.000363870 (0.000325906) 	 BC Loss: 0.000211202 (0.000164208)
Iteration: [ 47001/ 50000] 	 Loss: 0.001514807 (0.001019335) 	 Physics Loss: 0.000957777 (0.000599596) 	 Data Loss: 0.000467955 (0.000307101) 	 BC Loss: 0.000089074 (0.000112639)
Iteration: [ 48001/ 50000] 	 Loss: 0.000714825 (0.000956614) 	 Physics Loss: 0.000394674 (0.000525267) 	 Data Loss: 0.000164629 (0.000315504) 	 BC Loss: 0.000155522 (0.000115843)
Iteration: [ 49001/ 50000] 	 Loss: 0.000784852 (0.001260801) 	 Physics Loss: 0.000456999 (0.000808780) 	 Data Loss: 0.000215380 (0.000316373) 	 BC Loss: 0.000112472 (0.000135648)

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.5
Commit 760b2e5b739 (2025-04-14 06:53 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