lpe.livemd

Language Physics Engine

System.put_env("XLA_BUILD", "true")

Mix.install([
  {:nx, "~> 0.11.0"},
  {:exla, "~> 0.11.0"}
])

System.put_env("HSA_OVERRIDE_GFX_VERSION", "11.0.0")
System.put_env("XLA_TARGET", "rocm")

# ⚡ 1. Disable the buggy AMD Fusion pass
System.put_env("TF_ROCM_FUSION_ENABLE", "0")

# ⚡ 2. Disable XLA Autotuning completely
System.put_env("XLA_FLAGS", "--xla_gpu_autotune_level=0 --xla_gpu_enable_reassociation_for_converted_ar=false --xla_gpu_graph_enable_concurrent_region=false")

System.put_env("ROCM_DISABLE_CU_MEMSET", "1")

Application.put_env(:exla, :clients,
  rocm: [
    platform: :rocm,
    memory_fraction: 0.75,  # ← force total claimed under 8GB
    preallocate: false,
    num_replicas: 1
  ]
)

System.put_env("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.33")

Application.put_env(:exla, :default_client, :rocm)

gpu = {EXLA.Backend, client: :rocm}
Nx.global_default_backend(gpu)
Nx.Defn.default_options(compiler: EXLA)

platforms = EXLA.Client.get_supported_platforms()
if platforms[:rocm] do
  IO.puts("✅ AMD RX 7600 online!")
else
  IO.puts("❌ GPU missing — check HSA_OVERRIDE_GFX_VERSION")
end

LPSSM, Positional, Muon_Optimizer

# ===========================================================================
# THE TINY SSM (Stateful Predictor)
# ===========================================================================
defmodule DLM.TinySSM do
  import Nx.Defn

  defn gelu(x) do
    cdf = Nx.multiply(
      Nx.tensor(0.7978845608, type: :f32),
      Nx.add(x, Nx.multiply(Nx.tensor(0.044715, type: :f32), Nx.pow(x, 3)))
    )
    Nx.multiply(Nx.multiply(Nx.tensor(0.5, type: :f32), x), Nx.add(Nx.tensor(1.0, type: :f32), Nx.tanh(cdf)))
  end

  defn apply_svd(x, u, v) do
    hidden = Nx.dot(x, u)
    activated = gelu(hidden) 
    Nx.dot(activated, v)
  end

  defn softplus(x) do
    Nx.select(x > 20.0, x, Nx.log1p(Nx.exp(x)))
  end

  defn forward(x_t, h_prev, params) do
    # ⚡ Unpack the 7 parameters
    {a, bu, bv, cu, cv, du, dv} = params

    b_proj = apply_svd(x_t, bu, bv)
    
    # ⚡ The Dynamic Gate
    gate_proj = apply_svd(x_t, du, dv)
    delta = Nx.max(softplus(gate_proj), 1.0e-4) # Minimum step size for stability
    
    # Dynamic Friction & Force
    a_bar = Nx.exp(Nx.multiply(delta, -Nx.abs(a)))
    b_bar = Nx.multiply(delta, b_proj)
    
    h_new = Nx.tanh(Nx.add(Nx.multiply(a_bar, h_prev), b_bar))
    
    predicted_vector = apply_svd(h_new, cu, cv)
    
    {predicted_vector, h_new}
  end
end

{:module, DLM.TinySSM, <<70, 79, 82, 49, 0, 0, 23, ...>>, true}

defmodule DLM.Positional do
  def build(len, dim) do
    i = Nx.iota({len, 1})
    d = Nx.iota({1, dim})
    dim_f = Nx.tensor(dim, type: :f32) # 🔥 Force float math
    Nx.sin(Nx.divide(i, Nx.pow(10000.0, Nx.divide(d, dim_f))))
  end
end

{:module, DLM.Positional, <<70, 79, 82, 49, 0, 0, 8, ...>>, {:build, 2}}

# ===========================================================================
# HYBRID OPTIMIZER (AdamW for 1D, Muon for 2D)
# ===========================================================================
defmodule DLM.HybridMuon7 do
  import Nx.Defn

  # -------------------------------------------------------------------------
  # AdamW (For 1D Friction/Inertia vectors)
  # -------------------------------------------------------------------------
  defn adam_step(p, g, m, v, step, lr) do
    beta1 = 0.9
    beta2 = 0.999
    weight_decay = 0.01

    m_new = Nx.add(Nx.multiply(beta1, m), Nx.multiply(1.0 - beta1, g))
    v_new = Nx.add(Nx.multiply(beta2, v), Nx.multiply(1.0 - beta2, Nx.pow(g, 2)))

    m_hat = Nx.divide(m_new, Nx.subtract(1.0, Nx.pow(beta1, step)))
    v_hat = Nx.divide(v_new, Nx.subtract(1.0, Nx.pow(beta2, step)))

    update_val = Nx.multiply(lr, Nx.divide(m_hat, Nx.add(Nx.sqrt(v_hat), 1.0e-8)))

    p_decayed = Nx.subtract(p, Nx.multiply(Nx.multiply(lr, weight_decay), p))
    p_new = Nx.subtract(p_decayed, update_val)

    {p_new, m_new, v_new}
  end

  # -------------------------------------------------------------------------
  # Muon Orthogonalization (Newton-Schulz Iteration)
  # -------------------------------------------------------------------------
  defn newton_schulz(g) do
    # 1. Normalize the gradient matrix to stabilize the iteration
    norm = Nx.sqrt(Nx.sum(Nx.pow(g, 2)))
    x_init = Nx.divide(g, Nx.add(norm, 1.0e-8))

    # 2. Run 5 loops of: X = 1.5*X - 0.5*X*(X^T*X)
    {x_final, _} =
      while {x = x_init, i = 0}, Nx.less(i, 5) do
        x_t = Nx.transpose(x)
        x_t_x = Nx.dot(x_t, x)
        term = Nx.dot(x, x_t_x)
        
        x_next = Nx.subtract(Nx.multiply(1.5, x), Nx.multiply(0.5, term))
        {x_next, i + 1}
      end
      
    x_final
  end

  # -------------------------------------------------------------------------
  # Muon Step (For 2D Transformation Matrices)
  # -------------------------------------------------------------------------
  defn muon_step(p, g, m, v, lr) do
    beta1 = 0.95 # Muon standard momentum
    weight_decay = 0.01
    
    # Update Momentum
    m_new = Nx.add(Nx.multiply(beta1, m), Nx.multiply(1.0 - beta1, g))
    
    # ⚡ THE MAGIC: Orthogonalize the momentum!
    ortho_update = newton_schulz(m_new)
    
    # ⚡ THE FIX: Keep it entirely in the Nx Graph
    r = Nx.axis_size(p, 0)
    c = Nx.axis_size(p, 1)
    scale = Nx.max(r, c) |> Nx.as_type(:f32)
    
    scaled_update = Nx.multiply(ortho_update, scale)

    # Apply weight decay and step
    p_decayed = Nx.subtract(p, Nx.multiply(Nx.multiply(lr, weight_decay), p))
    p_new = Nx.subtract(p_decayed, Nx.multiply(lr, scaled_update))
    
    # Note: Muon doesn't use 'v' (variance), we just pass it back to keep the tuple shapes identical
    {p_new, m_new, v}
  end

  # -------------------------------------------------------------------------
  # The 7-Parameter Router
  # -------------------------------------------------------------------------
  defn update(params, grads, ms, vs, step, lr) do
    {a, bu, bv, cu, cv, du, dv} = params
    {ga, gbu, gbv, gcu, gcv, gdu, gdv} = grads
    {ma, mbu, mbv, mcu, mcv, mdu, mdv} = ms
    {va, vbu, vbv, vcu, vcv, vdu, vdv} = vs

    new_step = Nx.add(step, 1)

    # ⚡ 1D Scalar: AdamW
    {a_new, ma_new, va_new} = adam_step(a, ga, ma, va, new_step, lr)
    
    # ⚡ 2D Matrices: Muon Orthogonal Rotations
    {bu_new, mbu_new, vbu_new} = muon_step(bu, gbu, mbu, vbu, lr)
    {bv_new, mbv_new, vbv_new} = muon_step(bv, gbv, mbv, vbv, lr)
    {cu_new, mcu_new, vcu_new} = muon_step(cu, gcu, mcu, vcu, lr)
    {cv_new, mcv_new, vcv_new} = muon_step(cv, gcv, mcv, vcv, lr)
    {du_new, mdu_new, vdu_new} = muon_step(du, gdu, mdu, vdu, lr)
    {dv_new, mdv_new, vdv_new} = muon_step(dv, gdv, mdv, vdv, lr)

    new_params = {a_new, bu_new, bv_new, cu_new, cv_new, du_new, dv_new}
    new_ms = {ma_new, mbu_new, mbv_new, mcu_new, mcv_new, mdu_new, mdv_new}
    new_vs = {va_new, vbu_new, vbv_new, vcu_new, vcv_new, vdu_new, vdv_new}

    {new_params, new_ms, new_vs, new_step}
  end
end

{:module, DLM.HybridMuon7, <<70, 79, 82, 49, 0, 0, 38, ...>>, true}

defmodule DLM.PhysicsGenesis do
  import Nx.Defn

  @dim 512
  @rank 256

  defn ignite(key) do
    # a_diag: Friction/Inertia
    {a, k1} = Nx.Random.normal(key, -0.1, 0.01, shape: {@dim})
    
    # b: Input Force Matrices
    {bu, k2} = Nx.Random.normal(k1, 0.0, 0.02, shape: {@dim, @rank})
    {bv, k3} = Nx.Random.normal(k2, 0.0, 0.02, shape: {@rank, @dim})
    
    # c: Output Velocity Matrices
    {cu, k4} = Nx.Random.normal(k3, 0.0, 0.02, shape: {@dim, @rank})
    {cv, k5} = Nx.Random.normal(k4, 0.0, 0.02, shape: {@rank, @dim})

    # ⚡ NEW: d: Dynamic Gating Matrices (Initialized near zero)
    {du, k6} = Nx.Random.normal(k5, 0.0, 0.001, shape: {@dim, @rank})
    {dv, next_key} = Nx.Random.normal(k6, 0.0, 0.001, shape: {@rank, @dim})

    {{a, bu, bv, cu, cv, du, dv}, next_key}
  end

  def init_optimizer(params) do
    zero_state = DLM.Tree.map(params, fn p -> Nx.broadcast(0.0, Nx.shape(p)) end)
    {zero_state, zero_state, Nx.tensor(0)}
  end
end

warning: DLM.Tree.map/2 is undefined (module DLM.Tree is not available or is yet to be defined)
└─ lpe.livemd#cell:ix2bsm2bfbaikqjk:27: DLM.PhysicsGenesis.init_optimizer/1

{:module, DLM.PhysicsGenesis, <<70, 79, 82, 49, 0, 0, 16, ...>>, {:init_optimizer, 1}}

Inference

defmodule DLM.Tree do
  @moduledoc """
  Recursively maps a function over nested tuples of tensors.
  Works perfectly both in standard Elixir AND inside Nx `defn`.
  """
  
  def map(tuple, func) when is_tuple(tuple) do
    tuple
    |> Tuple.to_list()
    |> Enum.map(&map(&1, func))
    |> List.to_tuple()
  end

  # The base case: When it hits a tensor (or a compile-time Expr), apply the function
  def map(tensor, func) do
    func.(tensor)
  end
end

{:module, DLM.Tree, <<70, 79, 82, 49, 0, 0, 8, ...>>, {:map, 2}}

defmodule DLM.ContinuousProbe do
  import Nx.Defn

  # ⚡ Find the closest Gravity Well (Token)
  defn decode_coordinate(predicted_coord, all_token_coords) do
    # predicted_coord shape: {512}
    # all_token_coords shape: {256, 512}
    
    # Calculate Euclidean (squared) distance to every possible ASCII token
    distances = Nx.sum(Nx.pow(Nx.subtract(all_token_coords, predicted_coord), 2), axes: [-1])
    
    # Return the index of the closest coordinate (which maps 1:1 to the ASCII integer!)
    Nx.argmin(distances)
  end
end

{:module, DLM.ContinuousProbe, <<70, 79, 82, 49, 0, 0, 10, ...>>, true}

# ==============================================================================
# THE CONTINUOUS PROBE (SHAPE-CORRECTED)
# ==============================================================================
# ⚡ UPDATE THESE FILENAMES TO MATCH YOUR SAVED WEIGHTS
teacher_params = File.read!("v32_mse_geo_enc_4_checkpoint_step_1000.bin") |> :erlang.binary_to_term()
IO.puts("🧊 Loading Frozen Micro-Diffusion Tokenizer...")
{fw1, fw2} = File.read!("semantic_bytes_frozen_9_encoder.bin") |> :erlang.binary_to_term()
frozen_w1 = Nx.backend_copy(fw1, gpu)
frozen_w2 = Nx.backend_copy(fw2, gpu)

# create byte look up table
feature_table = DLM.CharFeaturizer.build_feature_table(range: 0..127)
gpu_table = Nx.backend_copy(feature_table, gpu)

# 🌌 Materialize the entire 512D Gravity Well Map
all_ascii = Nx.tensor(Enum.to_list(0..255), type: :u32) |> Nx.reshape({256, 1})
all_token_coords_raw = DLM.PhysicsEngine.get_coordinates(all_ascii, frozen_w1, frozen_w2, gpu_table)

# Extract the flat {256, 512} lookup table
all_token_coords = all_token_coords_raw[[.., 0, ..]] |> Nx.backend_copy(gpu)

# ⚡ THE FIX: Initialize Teacher's hidden state as a FLAT {512} vector
h_zeros = Nx.broadcast(0.0, {512}) |> Nx.backend_copy(gpu)

prompt_text = "Hello, my na"
prompt_tokens_tensor = DLM.LatentTokenizer.to_ascii_tokens(prompt_text)
prompt_tokens = Nx.to_flat_list(prompt_tokens_tensor)

IO.puts("🧠 Burning in the Pondering Teacher...")

# ⚡ BURN-IN
{final_h, last_coord} = 
  Enum.reduce(prompt_tokens, {h_zeros, nil}, fn tok, {h, _} ->
    # Get the perfect physical coordinate for this token
    tok_tensor = Nx.tensor([[tok]], type: :u32) |> Nx.backend_copy(gpu)
    actual_coord = DLM.PhysicsEngine.get_coordinates(tok_tensor, frozen_w1, frozen_w2, gpu_table)[0][0]
    
    # ⚡ The Teacher ponders the coordinate!
    {_, next_h} = DLM.PhysicsEngine.predict_step(actual_coord, h, teacher_params)
    
    {next_h, actual_coord}
  end)

IO.puts("\nPrompt: \"#{prompt_text}\"\n")
IO.write("Teacher: ")

# ⚡ CONTINUOUS GENERATION WITH TEMPERATURE
# Temperature (0.02 is a good start)
temp_scale = 0.05
initial_key = Nx.Random.key(42) |> Nx.backend_copy(gpu)

# ⚡ CONTINUOUS GENERATION WITH LANGEVIN TEMPERATURE
Enum.reduce(1..150, {final_h, last_coord, initial_key}, fn _step, {h, current_coord, cur_key} ->
  {predicted_coord, next_h} = DLM.PhysicsEngine.predict_step(current_coord, h, teacher_params)
  
  # ⚡ Langevin Kick
  {noise, next_key} = Nx.Random.normal(cur_key, 0.0, temp_scale, shape: {512})
  stochastic_coord = Nx.add(predicted_coord, noise)
  
  next_token_id = DLM.ContinuousProbe.decode_coordinate(stochastic_coord, all_token_coords)
  token_val = Nx.to_number(next_token_id)
  
  # 4. Telemetry (Void/Char logic)
  char_to_print = 
    cond do
      token_val in 32..126 -> <<token_val>>
      token_val in [9, 10, 13] -> <<token_val>> 
      true -> "░"
    end
    
  IO.write(char_to_print)
  
  # 5. Snap to the pure gravity well for the NEXT step
  # This prevents the noise from accumulating in the hidden state
  snapped_coord = all_token_coords[next_token_id]
  
  {next_h, snapped_coord, next_key}
end)
IO.puts("\n\n🐑 Dream Complete.")

🧊 Loading Frozen Micro-Diffusion Tokenizer...
🧠 Burning in the Pondering Teacher...

Prompt: "Hello, my na"

Teacher: hhhhhhhhhhhxhhhhxphxxhxhhhhhxhppphhhxhhhhhhhhxhhhxhxxhhhhhhhhhhhhxhxhhhxhhhhhhhhhhhxhhhhhhhhxhhhhxhhxhhhpxhhxhhhhhxhhhhxxhphhbhhhhhhhhhhhhhxhhhhhhhxhh

🐑 Dream Complete.

:ok

Stream

defmodule DLM.DataStreamer do
  def stream(path, seq_len, batch_size) do
    lines = path |> File.read!() |> String.split("\n", trim: true) |> Enum.shuffle()

    lines
    |> Enum.join("\n")
    |> to_charlist()
    |> Stream.chunk_every(seq_len + 1, seq_len, :discard)
    |> Stream.map(fn chunk ->
      ins = chunk |> Enum.take(seq_len) |> Nx.tensor(type: :s32)
      tgs = chunk |> Enum.drop(1)      |> Nx.tensor(type: :s32)
      {ins, tgs}
    end)
    |> Stream.chunk_every(batch_size, batch_size, :discard)
    |> Stream.map(fn batch ->
      {ins_list, tgs_list} = Enum.unzip(batch)
      {Nx.stack(ins_list), Nx.stack(tgs_list)}
    end)
  end

  @doc """
  Takes a massive 1D tensor of dataset tokens and creates an infinite stream
  of {batch_size, seq_len} matrices for the Physics Engine.
  """
  def build_infinite_stream(dataset_1d, batch_size, seq_len) do
    chunk_size = batch_size * seq_len
    total_tokens = Nx.size(dataset_1d)
    
    # Calculate how many full batches we can make
    total_batches = div(total_tokens, chunk_size)

    0..(total_batches - 1)
    |> Stream.cycle() # ⚡ Loops infinitely!
    |> Stream.map(fn batch_idx ->
      start_idx = batch_idx * chunk_size
      
      # Slice out the flat chunk and reshape it for the batch
      dataset_1d
      |> Nx.slice_along_axis(start_idx, chunk_size, axis: 0)
      |> Nx.reshape({batch_size, seq_len})
    end)
  end
end

{:module, DLM.DataStreamer, <<70, 79, 82, 49, 0, 0, 19, ...>>, {:build_infinite_stream, 3}}

defmodule DLM.CharFeaturizer do
  import Bitwise

  # UTF-8 encodes a codepoint into exactly 4 bytes, zero-padded
  # Follows the standard UTF-8 bit layout:
  #   1 byte:  0xxxxxxx
  #   2 bytes: 110xxxxx 10xxxxxx
  #   3 bytes: 1110xxxx 10xxxxxx 10xxxxxx
  #   4 bytes: 11110xxx 10xxxxxx 10xxxxxx 10xxxxxx
  def codepoint_to_utf8_bytes(cp) when cp <= 0x7F do
    [cp, 0, 0, 0]
  end

  def codepoint_to_utf8_bytes(cp) when cp <= 0x7FF do
    b1 = 0b11000000 ||| (cp >>> 6)
    b2 = 0b10000000 ||| (cp &&& 0b00111111)
    [b1, b2, 0, 0]
  end

  def codepoint_to_utf8_bytes(cp) when cp <= 0xFFFF do
    b1 = 0b11100000 ||| (cp >>> 12)
    b2 = 0b10000000 ||| ((cp >>> 6) &&& 0b00111111)
    b3 = 0b10000000 ||| (cp &&& 0b00111111)
    [b1, b2, b3, 0]
  end

  def codepoint_to_utf8_bytes(cp) do
    b1 = 0b11110000 ||| (cp >>> 18)
    b2 = 0b10000000 ||| ((cp >>> 12) &&& 0b00111111)
    b3 = 0b10000000 ||| ((cp >>> 6) &&& 0b00111111)
    b4 = 0b10000000 ||| (cp &&& 0b00111111)
    [b1, b2, b3, b4]
  end

  # Extracts all 32 bits from a 4-byte list, MSB first, as floats
  # [0x41, 0, 0, 0] → [0,1,0,0,0,0,0,1, 0,0,...,0]  (32 values)
  def extract_bits(bytes) do
    bytes
    |> Enum.flat_map(fn byte ->
      Enum.map(7..0//-1, fn shift ->
        (byte >>> shift) &&& 1
      end)
    end)
  end

  def unicode_flags(cp) do
    is_digit     = if cp in ?0..?9,            do: 1.0, else: 0.0
    is_lower     = if cp in ?a..?z,            do: 1.0, else: 0.0
    is_upper     = if cp in ?A..?Z,            do: 1.0, else: 0.0
    is_vowel     = if cp in ~c"aeiouAEIOU",      do: 1.0, else: 0.0
    is_space     = if cp in [?\s, ?\t, ?\n],   do: 1.0, else: 0.0
    is_punct     = if not (cp in ?0..?9) and
                      not (cp in ?a..?z) and
                      not (cp in ?A..?Z) and
                      not (cp in [?\s, ?\t, ?\n]) and
                      cp >= 32,                do: 1.0, else: 0.0
  
    # Scale by 4.0 so flags carry equal weight to the 32 byte bits
    [is_digit, is_lower, is_upper, is_vowel, is_space, is_punct]
    |> Enum.map(&(&1 * 1.5))
  end
  
  def build_feature_table(opts \\ []) do
    range = Keyword.get(opts, :range, 0..127)
    
    range
    |> Enum.map(fn cp ->
      utf8_bits = cp |> codepoint_to_utf8_bytes() |> extract_bits()
      flags     = unicode_flags(cp)
      utf8_bits ++ flags  # 38-dim total
    end)
    |> Nx.tensor(type: :f32)
  end
end

{:module, DLM.CharFeaturizer, <<70, 79, 82, 49, 0, 0, 24, ...>>, {:build_feature_table, 1}}

defmodule DLM.LatentTokenizer do
  import Nx.Defn

  # deftransform runs pure Elixir at compile time — safe to use Tuple.append here
  deftransform append_feat_dim(token_shape, feat_dim) do
    Tuple.insert_at(token_shape, tuple_size(token_shape), feat_dim)
  end

  defn to_utf8_bits(tokens, table) do
    feat_dim = Nx.axis_size(table, 1)
    flat = Nx.flatten(tokens)
    looked_up = Nx.take(table, flat)

    new_shape = append_feat_dim(Nx.shape(tokens), feat_dim)

    Nx.reshape(looked_up, new_shape)
  end

  @doc """
  Converts a standard Elixir string into a 1D tensor of ASCII integers.
  Example: "Cat" -> #Nx.Tensor<[67, 97, 116]>
  """
  def to_ascii_tokens(string) do
    string
    |> String.to_charlist()
    |> Nx.tensor(type: :u8) # u8 matches your 0-255 ASCII setup
  end

  defn softplus(x) do
    Nx.select(x > 20.0, x, Nx.log1p(Nx.exp(x)))
  end

  # ⚡ Add rms_norm to the tokenizer
  defn rms_norm(x, epsilon \\ 1.0e-6) do
    variance = Nx.mean(Nx.pow(x, 2), axes: [-1], keep_axes: true)
    x * Nx.rsqrt(variance + epsilon)
  end

  defn forward(tokens, params, table, key, noise_scale \\ 0.15) do
    {enc_w1, enc_w2, dec_w1, dec_w2} = params

    bits = to_utf8_bits(tokens, table)

    # Squeeze the middle dim: {batch, 1, 38} → {batch, 38}
    bits = Nx.squeeze(bits, axes: [1])
    
    hidden_enc = softplus(Nx.dot(bits, enc_w1))
    
    # ⚡ FIX: The Latent Trap! Normalize the vectors BEFORE adding noise.
    raw_latents = Nx.dot(hidden_enc, enc_w2)
    latents = rms_norm(raw_latents)

    {noise, next_key} = Nx.Random.normal(key, 0.0, 1.0, shape: Nx.shape(latents))
    noisy_latents = latents + (noise * noise_scale)

    hidden_dec = softplus(Nx.dot(noisy_latents, dec_w1))
    logits = Nx.dot(hidden_dec, dec_w2)

    {logits, noisy_latents, next_key}
  end
end

{:module, DLM.LatentTokenizer, <<70, 79, 82, 49, 0, 0, 33, ...>>, true}

defmodule DLM.LatentRadar do
  import Nx.Defn

  # 1. Recreate the exact deterministic path from your Tokenizer
  defn get_coordinates(tokens, enc_w1, enc_w2, table) do
    bits = DLM.LatentTokenizer.to_utf8_bits(tokens, table)
    hidden_enc = DLM.LatentTokenizer.softplus(Nx.dot(bits, enc_w1))
    raw_latents = Nx.dot(hidden_enc, enc_w2)
    
    # The crucial latents trap!
    DLM.LatentTokenizer.rms_norm(raw_latents) 
  end

  # 2. Calculate the physical distance between all 128 points simultaneously
  defn pairwise_distances(coords) do
    # Expand shapes to {128, 1, 512} and {1, 128, 512}
    a = Nx.new_axis(coords, 1)
    b = Nx.new_axis(coords, 0)
    
    # Calculate the difference, square it, sum along the 512D axis, and take the square root
    diff = Nx.subtract(a, b)
    Nx.sqrt(Nx.sum(Nx.pow(diff, 2), axes: [-1])) 
  end
end

{:module, DLM.LatentRadar, <<70, 79, 82, 49, 0, 0, 16, ...>>, true}

Train

defmodule DLM.Loss do
  import Nx.Defn

  defn cross_entropy(logits, targets) do
    vocab_size = Nx.axis_size(logits, -1)
    
    # Create a 1-hot vector for the correct answer
    one_hot = Nx.equal(Nx.iota({vocab_size}), Nx.new_axis(targets, -1))
    
    # Stable Log-Softmax
    max_logits = Nx.reduce_max(logits, axes: [-1], keep_axes: true)
    safe_logits = Nx.subtract(logits, max_logits)
    log_probs = Nx.subtract(safe_logits, Nx.log(Nx.sum(Nx.exp(safe_logits), axes: [-1], keep_axes: true)))
    
    # Cross Entropy calculation
    Nx.mean(Nx.negate(Nx.sum(Nx.multiply(one_hot, log_probs), axes: [-1])))
  end
end

{:module, DLM.Loss, <<70, 79, 82, 49, 0, 0, 12, ...>>, true}

defmodule DLM.SinkhornOT do
  import Nx.Defn

  defn cost_matrix(pred, target) do
    p_exp = Nx.new_axis(pred, 1)
    t_exp = Nx.new_axis(target, 0)
    diff = Nx.subtract(p_exp, t_exp)
    Nx.sum(Nx.pow(diff, 2), axes: [-1])
  end

  # ⚡ FIX 1: Bump the default epsilon up to 10.0 to handle your 512D spatial scale
  defn compute(pred, target, epsilon \\ 10.0, max_iters \\ 5) do
    c = cost_matrix(pred, target)
    k = Nx.exp(Nx.divide(Nx.negate(c), epsilon))
    
    batch_size = Nx.axis_size(pred, 0)
    mass = 1.0 / batch_size
    mu = Nx.broadcast(mass, {batch_size})
    nu = Nx.broadcast(mass, {batch_size})

    {_final_k, u_final, v_final, _mu, _nu, _i} =
      while {k_mat = k, 
             u = Nx.broadcast(1.0, {batch_size}), 
             _v = Nx.broadcast(1.0, {batch_size}), 
             m = mu, n = nu, i = 0},
            Nx.less(i, max_iters) do
            
        # ⚡ FIX 2: Add 1.0e-8 to the denominator to mathematically prevent NaN crashes
        denom_v = Nx.add(Nx.dot(Nx.transpose(k_mat), u), 1.0e-8)
        v_new = Nx.divide(n, denom_v)
        
        denom_u = Nx.add(Nx.dot(k_mat, v_new), 1.0e-8)
        u_new = Nx.divide(m, denom_u)

        {k_mat, u_new, v_new, m, n, i + 1}
      end

    u_exp = Nx.new_axis(u_final, 1)
    v_exp = Nx.new_axis(v_final, 0)
    p = Nx.multiply(Nx.multiply(u_exp, k), v_exp)

    Nx.sum(Nx.multiply(p, c))
  end
end

{:module, DLM.SinkhornOT, <<70, 79, 82, 49, 0, 0, 25, ...>>, true}

defmodule DLM.PhysicsLoss do
  import Nx.Defn

  # ⚡ Universal L2 Normalizer (Works on 2D, 3D, N-D tensors)
  defn l2_normalize(tensor) do
    # Calculate Euclidean norm across the last axis natively
    norm = Nx.sqrt(Nx.sum(Nx.pow(tensor, 2), axes: [-1], keep_axes: true))
    # Divide by norm with a tiny epsilon to prevent divide-by-zero
    Nx.divide(tensor, Nx.add(norm, 1.0e-6))
  end

  # ⚡ 1. Angular Alignment (Cosine Distance)
  defn cosine_distance(pred, target) do
    p_norm = l2_normalize(pred)
    t_norm = l2_normalize(target)

    # Dot product across the 512 dimension
    cos_sim = Nx.sum(Nx.multiply(p_norm, t_norm), axes: [-1])
    
    # Minimize distance (1.0 - similarity)
    Nx.mean(Nx.subtract(1.0, cos_sim))
  end

  # ⚡ 2. Repulsion (Memory-Safe Batched Soft-NN for Rank 2 Tensors)
  defn soft_nn_loss(pred, target, temperature \\ 0.1) do
    # pred and target are {batch_size, dim}
    p_l2 = l2_normalize(pred)
    t_l2 = l2_normalize(target)

    # 1. Batched Dot Product: {batch, dim} x {dim, batch} -> {batch, batch}
    # This creates a similarity matrix comparing every pred to every target in the batch
    sim_tensor = Nx.dot(p_l2, Nx.transpose(t_l2))
    sim_tensor = Nx.divide(sim_tensor, temperature)

    # 2. Create Identity Targets -> {batch, batch}
    batch_size = Nx.axis_size(pred, 0)
    # ⚡ THE FIX: Pass a 2D tuple {32, 32} instead of a 1D tuple {32}
    labels = Nx.eye({batch_size, batch_size}, type: :f32)

    # 3. Manual Dense Log-Softmax
    # Step A: Max trick for numerical stability
    max_sim = Nx.reduce_max(sim_tensor, axes: [-1], keep_axes: true)
    shifted_sim = Nx.subtract(sim_tensor, max_sim)
    
    # Step B: Log(Sum(Exp))ozen_w2,
    log_sum_exp = Nx.log(Nx.sum(Nx.exp(shifted_sim), axes: [-1], keep_axes: true))
    log_softmax = Nx.subtract(shifted_sim, log_sum_exp)

    # Step C: -Sum(Targets * Log_Softmax)
    Nx.multiply(labels, log_softmax)
    |> Nx.sum(axes: [-1])
    |> Nx.negate()
    |> Nx.mean()
  end

  # defn compute(pred, target) do
  #   mse = Nx.mean(Nx.pow(Nx.subtract(pred, target), 2))
  #   cos = cosine_distance(pred, target)
    
  #   # ⚡ Update the explicit epsilon call here to 10.0
  #   ot_loss = DLM.SinkhornOT.compute(pred, target, 5.0, 5)

  #   Nx.add(Nx.add(mse, cos), Nx.multiply(ot_loss, 0.5))
  # end

  defn compute(pred, target) do
    # Simple MSE in 512D coordinate space
    # The encoder's geometry makes this meaningful
    Nx.mean(Nx.pow(Nx.subtract(pred, target), 2))
  end
end

{:module, DLM.PhysicsLoss, <<70, 79, 82, 49, 0, 0, 25, ...>>, true}

# ===========================================================================
# THE PHYSICS ENGINE (Geometric Loss & BPTT)
# ===========================================================================
defmodule DLM.PhysicsEngine do
  import Nx.Defn

  defn get_coordinates(tokens, enc_w1, enc_w2, feature_table) do
  bits = DLM.LatentTokenizer.to_utf8_bits(tokens, feature_table)
  # bits: {batch, seq_len, 38}
  # Nx.dot over last axis: {batch, seq_len, 38} × {38, 256} → {batch, seq_len, 256}
  hidden_enc = DLM.LatentTokenizer.softplus(Nx.dot(bits, enc_w1))
  raw_latents = Nx.dot(hidden_enc, enc_w2)
  DLM.LatentTokenizer.rms_norm(raw_latents)
  # output: {batch, seq_len, 512} ✅
end

  # # ⚡ RECURRENT PREDICT STEP (Gradient-Safe Pondering)
  # defn predict_step(current_coord, h_prev, predictor_params) do
  #   max_loops = 5 
    
  #   # ⚡ We use the SQUARED epsilon to avoid the Nx.sqrt() derivative explosion!
  #   # If our target threshold is 0.1 distance, 0.1^2 = 0.01.
  #   sq_epsilon = Nx.tensor(0.01, type: :f32) 
    
  #   init_delta = Nx.tensor(100.0, type: :f32)
  #   init_pred = Nx.broadcast(0.0, Nx.shape(current_coord))

  #   {_k, _c, final_h, _d, _p, _m, _e, final_pred} = 
  #     while {k = 0, 
  #            c_coord = current_coord, 
  #            c_h = h_prev, 
  #            delta = init_delta, 
  #            p = predictor_params, 
  #            m = max_loops, 
  #            e = sq_epsilon,
  #            prev_pred = init_pred},
  #           Nx.logical_and(Nx.less(k, m), Nx.greater(delta, e)) do
            
  #       {new_pred, next_h} = DLM.TinySSM.forward(c_coord, c_h, p)
        
  #       change_in_thought = Nx.subtract(new_pred, prev_pred)
        
  #       # ⚡ NO SQUARE ROOT! We just use the raw sum of squares. 
  #       # This is computationally faster and 100% gradient-safe.
  #       token_movements_sq = Nx.sum(Nx.pow(change_in_thought, 2), axes: [-1])
        
  #       # Check the most confused token's SQUARED movement
  #       max_movement_sq = Nx.reduce_max(token_movements_sq)
        
  #       {k + 1, c_coord, next_h, max_movement_sq, p, m, e, new_pred}
  #     end

  #   # The final landing spot is the original coordinate + the final, fully reasoned jump
  #   final_coord = Nx.add(current_coord, final_pred)

  #   {final_coord, final_h}
  # end

  defn predict_step(current_coord, h_prev, predictor_params) do
    # Single step — clean gradient path
    {new_pred, next_h} = DLM.TinySSM.forward(current_coord, h_prev, predictor_params)
    # final_coord = Nx.add(current_coord, new_pred)
    {new_pred, next_h}
  end

  # ⚡ Added noise_scale to the arguments
  defn compute_sequence_loss(input_coords, target_coords, h_init, predictor_params, noise_key, noise_scale) do
    {batch_size, seq_len, dim} = Nx.shape(input_coords)
    inputs_time_first = Nx.transpose(input_coords, axes: [1, 0, 2])
    targets_time_first = Nx.transpose(target_coords, axes: [1, 0, 2])

    result =
      while {i = 0, total_loss = Nx.tensor(0.0, type: :f32), h = h_init, 
             inputs = inputs_time_first, targets = targets_time_first, p = predictor_params,
             key = noise_key, ns = noise_scale},
            Nx.less(i, seq_len) do
      
        # 1. Generate Kinetic Noise
        {noise, next_key} = Nx.Random.normal(key, 0.0, 1.0, shape: {batch_size, dim})
        
        # ⚡ 2. Multiply by the dynamic noise_scale (0.0 means noise is completely zeroed out)
        noisy_input = Nx.add(inputs[i], Nx.multiply(noise, ns))
      
        # 3. Physics step and Loss calculation
        {final_coord, h_new} = predict_step(noisy_input, h, p)
        step_loss = DLM.PhysicsLoss.compute(final_coord, targets[i])
        
        {i + 1, Nx.add(total_loss, step_loss), h_new, inputs, targets, p, next_key, ns}
      end

    accumulated_loss = elem(result, 1)
    h_final = elem(result, 2)
    {Nx.divide(accumulated_loss, seq_len), h_final}
  end

  defn compute_grad(params, batch_tokens, h_init, frozen_w1, frozen_w2, noise_key, noise_scale, gpu_table) do
    {_batch_size, full_seq_len} = Nx.shape(batch_tokens)
    seq_len = full_seq_len - 1
    
    input_tokens  = Nx.slice_along_axis(batch_tokens, 0, seq_len, axis: 1)
    target_tokens = Nx.slice_along_axis(batch_tokens, 1, seq_len, axis: 1)
  
    input_coords  = get_coordinates(input_tokens,  frozen_w1, frozen_w2, gpu_table)
    target_coords = get_coordinates(target_tokens, frozen_w1, frozen_w2, gpu_table)
  
    # ✅ Call 1: get gradients — value_and_grad only sees the scalar loss
    {loss, raw_grads} = value_and_grad(params, fn p ->
      {seq_loss, _h_f} = compute_sequence_loss(input_coords, target_coords, h_init, p, noise_key, noise_scale)
      seq_loss
    end)
  
    # ✅ Call 2: get h_final — plain call outside value_and_grad
    {_loss, h_final} = compute_sequence_loss(input_coords, target_coords, h_init, params, noise_key, noise_scale)
  
    {g_a, g_bu, g_bv, g_cu, g_cv, g_du, g_dv} = raw_grads
  
    clipped_grads = {
      Nx.clip(g_a,   -1.0, 1.0),
      Nx.clip(g_bu,  -1.0, 1.0),
      Nx.clip(g_bv,  -1.0, 1.0),
      Nx.clip(g_cu,  -1.0, 1.0),
      Nx.clip(g_cv,  -1.0, 1.0),
      Nx.clip(g_du,  -1.0, 1.0),
      Nx.clip(g_dv,  -1.0, 1.0)
    }
  
    {loss, h_final, clipped_grads}
  end
end

{:module, DLM.PhysicsEngine, <<70, 79, 82, 49, 0, 0, 40, ...>>, true}

defmodule DLM.CheckpointManager do
  @moduledoc """
  Scans the directory for the latest checkpoint of a specific run prefix.
  If no checkpoints exist, it triggers a completely fresh initialization.
  """
  def get_resume_state(run_prefix) do
    checkpoints = Path.wildcard("#{run_prefix}_checkpoint_step_*.bin")

    if Enum.empty?(checkpoints) do
      if File.exists?("#{run_prefix}_final.bin") do
        IO.puts("✨ Found completed run: #{run_prefix}_final.bin")
        {"#{run_prefix}_final.bin", :completed}
      else
        IO.puts("✨ No checkpoints found for '#{run_prefix}'. Starting FRESH from Genesis.")
        {:fresh, 0}
      end
    else
      # Extract step numbers using regex and find the maximum
      latest_file = Enum.max_by(checkpoints, fn file ->
        case Regex.run(~r/_step_(\d+)\.bin$/, file) do
          [_, step_str] -> String.to_integer(step_str)
          _ -> -1
        end
      end)

      [_, step_str] = Regex.run(~r/_step_(\d+)\.bin$/, latest_file)
      step = String.to_integer(step_str)
      
      IO.puts("🚀 Auto-Resuming '#{run_prefix}' from Step #{step}...")
      {latest_file, step}
    end
  end
end

{:module, DLM.CheckpointManager, <<70, 79, 82, 49, 0, 0, 17, ...>>, {:get_resume_state, 1}}

# ===========================================================================
# THE EXECUTION LOOP (With Auto-Resume & Fresh Genesis logic)
# ===========================================================================
# ⚡ 1. DEFINE YOUR RUN CONFIGURATION HERE
run_prefix  = "v32_mse_geo_enc_4"  
total_steps = 15000

# ⚡ 2. CHECK RESUME STATE (No more base_weights passed in!)
{load_path, start_step} = DLM.CheckpointManager.get_resume_state(run_prefix)

if start_step == :completed do
  IO.puts("✅ This run has already reached #{total_steps} steps. Exiting.")
else
  IO.puts("🧊 Loading Frozen Micro-Diffusion Tokenizer...")
  {fw1, fw2} = File.read!("semantic_bytes_frozen_9_encoder.bin") |> :erlang.binary_to_term()
  
  gpu = {EXLA.Backend, client: :rocm}
  frozen_w1 = Nx.backend_copy(fw1, gpu)
  frozen_w2 = Nx.backend_copy(fw2, gpu)

  batch_size = 32
  seq_len = 128  

  IO.puts("🧠 Loading LPSSM 128seqlen training data CPU Memory...")
  dataset_2d = File.read!("scholar_full_mixed_128.bin") |> :erlang.binary_to_term()
  dataset = Nx.flatten(dataset_2d)
  data_stream = DLM.DataStreamer.build_infinite_stream(dataset, batch_size, seq_len)

  # ⚡ 3. BRANCHING WEIGHT INITIALIZATION
  {starting_params, initial_key} =
    if load_path == :fresh do
      IO.puts("✨ Igniting brand new Genesis weights for #{run_prefix}...")
      DLM.PhysicsGenesis.ignite(Nx.Random.key(System.system_time()))
    else
      IO.puts("✨ Loading Weights from #{load_path}...")
      loaded_p = File.read!(load_path) |> :erlang.binary_to_term()
      # Give it a fresh noise key even if resuming
      {loaded_p, Nx.Random.key(System.system_time())}
    end

  # Initialize the optimizer states with zeroes
  {m, v, step_count} = DLM.PhysicsGenesis.init_optimizer(starting_params)

  # Move everything to the GPU
  p = DLM.Tree.map(starting_params, &Nx.backend_copy(&1, gpu))
  m = DLM.Tree.map(m, &Nx.backend_copy(&1, gpu))
  v = DLM.Tree.map(v, &Nx.backend_copy(&1, gpu))
  step_count = Nx.tensor(start_step) |> Nx.backend_copy(gpu)

  # create byte look up table
  feature_table = DLM.CharFeaturizer.build_feature_table(range: 0..127)
  gpu_table = Nx.backend_copy(feature_table, gpu)

  # ⚡ 4. ZIP FROM THE START_STEP
  {final_p, _m, _v, _key, _step} = 
    Enum.reduce(Enum.zip(start_step..(total_steps - 1), data_stream), {p, m, v, initial_key, step_count}, fn {step, cpu_batch}, {cp, cm, cv, cur_key, c_step} ->
      
      gpu_batch = Nx.backend_copy(cpu_batch, gpu)
      
      # ⚡ DYNAMIC LEARNING RATE SCHEDULER
      max_lr = 1.2e-5     # Enough power to reach the canyon, slow enough to stay in it
      warmup_steps = 800  # Give the fluid more time to build momentum

      current_lr =
        if step < warmup_steps do
          # Linear Warmup: Gently ramp up to max_lr
          max_lr * (step / warmup_steps)
        else
          # Cosine Decay: Smoothly brake as we approach total_steps
          decay_ratio = (step - warmup_steps) / (total_steps - warmup_steps)
          max_lr * 0.5 * (1.0 + :math.cos(:math.pi() * decay_ratio))
        end
      
      # ⚡ Inject a gentle 5% Brownian motion to prevent limit cycles
      noise_scale = 0.05 # Change to 0.02 when testing the Drunk Driver
      
      lr_tensor = Nx.tensor(current_lr, type: :f32) |> Nx.backend_copy(gpu)
      noise_tensor = Nx.tensor(noise_scale, type: :f32) |> Nx.backend_copy(gpu)
      h_init = Nx.broadcast(0.0, {batch_size, 512}) |> Nx.as_type(:f32) |> Nx.backend_copy(gpu)

      split_keys = Nx.Random.split(cur_key)
      noise_key = split_keys[0]
      next_key = split_keys[1]
      
      {loss_tensor, h_final, grads} = DLM.PhysicsEngine.compute_grad(
        cp, gpu_batch, h_init, frozen_w1, frozen_w2, noise_key, noise_tensor, gpu_table
      )

      # ⚡ MEMORY FIX 1: Read the loss into Elixir BEFORE the optimizer mutates memory
      loss_val = if rem(step, 50) == 0, do: Nx.to_number(loss_tensor), else: 0.0

      {up_p, up_m, up_v, up_step} = DLM.HybridMuon7.update(cp, grads, cm, cv, c_step, lr_tensor)

      # ➡️ Telemetry
      if rem(step, 50) == 0 do
        {_a, _bu, bv, _cu, _cv, _du, dv} = up_p
        h_norm = h_final |> Nx.pow(2) |> Nx.mean() |> Nx.sqrt() |> Nx.to_number()
        bv_norm = bv |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number()
        dv_norm = dv |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number()

        IO.puts("➡️  Step #{step} | Geometric Loss: #{Float.round(loss_val, 4)} | LR: #{Float.round(current_lr, 6)} | bv: #{Float.round(bv_norm, 4)} | dv: #{Float.round(dv_norm, 4)} | h_norm: #{h_norm}")
      end

      # 💾 DYNAMIC CHECKPOINT SAVING
      if rem(step, 1000) == 0 and step > 0 do
        IO.puts("💾 Checkpoint reached! Saving Step #{step}...")
        
        # ⚡ MEMORY FIX 2: Use backend_copy instead of backend_transfer mid-loop!
        cpu_snapshot = DLM.Tree.map(up_p, &Nx.backend_copy(&1, Nx.BinaryBackend))
        filename = "#{run_prefix}_checkpoint_step_#{step}.bin"
        File.write!(filename, :erlang.term_to_binary(cpu_snapshot))
        
        :erlang.garbage_collect()
      end

      {up_p, up_m, up_v, next_key, up_step}
    end)

  IO.puts("💾 Saving Final Gated Physics Engine Weights...")
  # It is safe to use backend_transfer here because the run is over
  good_params = DLM.Tree.map(final_p, &Nx.backend_transfer(&1, Nx.BinaryBackend))
  File.write!("#{run_prefix}_final.bin", :erlang.term_to_binary(good_params))
  IO.puts("✅ Resaved with clean CPU tensors")
  IO.puts("✅ Physics Engine Training Complete!")
end

warning: variable "step_count" is unused (if the variable is not meant to be used, prefix it with an underscore)
└─ lpe.livemd#cell:ztzjgm45vdvpmyur:42

✨ No checkpoints found for 'v32_mse_geo_enc_4'. Starting FRESH from Genesis.
🧊 Loading Frozen Micro-Diffusion Tokenizer...
🧠 Loading LPSSM 128seqlen training data CPU Memory...
✨ Igniting brand new Genesis weights for v32_mse_geo_enc_4...

17:55:25.846 [info] Merging Dots in computation: region_2.12

17:55:25.846 [info] Merging Dots in computation: region_24.36.clone

17:55:25.846 [info] Merging Dots in computation: region_13.23.clone.clone

17:55:25.846 [info] Merging Dots in computation: main.38

17:56:46.735 [info] Merging Dots in computation: region_1.2

17:56:46.735 [info] Merging Dots in computation: region_4.5

17:56:46.735 [info] Merging Dots in computation: region_7.8

17:56:46.735 [info] Merging Dots in computation: region_10.11

17:56:46.735 [info] Merging Dots in computation: region_13.14

17:56:46.735 [info] Merging Dots in computation: region_16.17
➡️  Step 0 | Geometric Loss: 0.9983 | LR: 0.0 | bv: 7.2311 | dv: 0.3612 | h_norm: 0.3781622350215912
➡️  Step 50 | Geometric Loss: 0.9938 | LR: 1.0e-6 | bv: 7.2311 | dv: 0.3614 | h_norm: 0.38016435503959656
➡️  Step 100 | Geometric Loss: 0.9802 | LR: 2.0e-6 | bv: 7.2313 | dv: 0.3642 | h_norm: 0.39308369159698486
➡️  Step 150 | Geometric Loss: 0.9545 | LR: 2.0e-6 | bv: 7.232 | dv: 0.3735 | h_norm: 0.4282442629337311
➡️  Step 200 | Geometric Loss: 0.9062 | LR: 3.0e-6 | bv: 7.2339 | dv: 0.3938 | h_norm: 0.4901507496833801
➡️  Step 250 | Geometric Loss: 0.8193 | LR: 4.0e-6 | bv: 7.2374 | dv: 0.4305 | h_norm: 0.5732743144035339
➡️  Step 300 | Geometric Loss: 0.68 | LR: 5.0e-6 | bv: 7.2436 | dv: 0.4879 | h_norm: 0.6662330031394958
➡️  Step 350 | Geometric Loss: 0.4874 | LR: 5.0e-6 | bv: 7.2536 | dv: 0.5692 | h_norm: 0.7609021663665771
➡️  Step 400 | Geometric Loss: 0.2692 | LR: 6.0e-6 | bv: 7.2679 | dv: 0.6739 | h_norm: 0.846370279788971
➡️  Step 450 | Geometric Loss: 0.0949 | LR: 7.0e-6 | bv: 7.2874 | dv: 0.7981 | h_norm: 0.9174398183822632
➡️  Step 500 | Geometric Loss: 0.0258 | LR: 8.0e-6 | bv: 7.3075 | dv: 0.9083 | h_norm: 0.9556823372840881
➡️  Step 550 | Geometric Loss: 0.0208 | LR: 8.0e-6 | bv: 7.3155 | dv: 0.9664 | h_norm: 0.9654875993728638
➡️  Step 600 | Geometric Loss: 0.0193 | LR: 9.0e-6 | bv: 7.3402 | dv: 1.0939 | h_norm: 0.9869988560676575
➡️  Step 650 | Geometric Loss: 0.019 | LR: 1.0e-5 | bv: 7.3771 | dv: 1.2199 | h_norm: 0.9947815537452698
➡️  Step 700 | Geometric Loss: 0.019 | LR: 1.1e-5 | bv: 7.3845 | dv: 1.2947 | h_norm: 0.9964059591293335
➡️  Step 750 | Geometric Loss: 0.0195 | LR: 1.1e-5 | bv: 7.3883 | dv: 1.3595 | h_norm: 0.9938599467277527
➡️  Step 800 | Geometric Loss: 0.0193 | LR: 1.2e-5 | bv: 7.3956 | dv: 1.4176 | h_norm: 0.9908983707427979
➡️  Step 850 | Geometric Loss: 0.0194 | LR: 1.2e-5 | bv: 7.4052 | dv: 1.4791 | h_norm: 0.9907137751579285
➡️  Step 900 | Geometric Loss: 0.0191 | LR: 1.2e-5 | bv: 7.4199 | dv: 1.5293 | h_norm: 0.9900738000869751
➡️  Step 950 | Geometric Loss: 0.0193 | LR: 1.2e-5 | bv: 7.4326 | dv: 1.5824 | h_norm: 0.9906759858131409
➡️  Step 1000 | Geometric Loss: 0.0193 | LR: 1.2e-5 | bv: 7.4505 | dv: 1.6453 | h_norm: 0.9901159405708313
💾 Checkpoint reached! Saving Step 1000...
➡️  Step 1050 | Geometric Loss: 0.0189 | LR: 1.2e-5 | bv: 7.4692 | dv: 1.6869 | h_norm: 0.9900822639465332
➡️  Step 1100 | Geometric Loss: 0.0182 | LR: 1.2e-5 | bv: 7.4981 | dv: 1.7616 | h_norm: 0.9908345341682434
➡️  Step 1150 | Geometric Loss: 0.018 | LR: 1.2e-5 | bv: 7.5266 | dv: 1.8682 | h_norm: 0.9897127747535706
➡️  Step 1200 | Geometric Loss: 0.0183 | LR: 1.2e-5 | bv: 7.5543 | dv: 1.9477 | h_norm: 0.9844520688056946
➡️  Step 1250 | Geometric Loss: 0.0181 | LR: 1.2e-5 | bv: 7.5814 | dv: 2.0104 | h_norm: 0.9825668931007385
➡️  Step 1300 | Geometric Loss: 0.018 | LR: 1.2e-5 | bv: 7.6182 | dv: 2.0798 | h_norm: 0.9797338843345642
➡️  Step 1350 | Geometric Loss: 0.0179 | LR: 1.2e-5 | bv: 7.6574 | dv: 2.1645 | h_norm: 0.9729112982749939
➡️  Step 1400 | Geometric Loss: 0.0179 | LR: 1.2e-5 | bv: 7.682 | dv: 2.2168 | h_norm: 0.9704632759094238
➡️  Step 1450 | Geometric Loss: 0.0174 | LR: 1.2e-5 | bv: 7.6968 | dv: 2.268 | h_norm: 0.962762176990509
➡️  Step 1500 | Geometric Loss: 0.0177 | LR: 1.2e-5 | bv: 7.7112 | dv: 2.3151 | h_norm: 0.9606000781059265
➡️  Step 1550 | Geometric Loss: 0.0171 | LR: 1.2e-5 | bv: 7.7288 | dv: 2.3668 | h_norm: 0.9601413011550903
➡️  Step 1600 | Geometric Loss: 0.0163 | LR: 1.2e-5 | bv: 7.745 | dv: 2.4013 | h_norm: 0.9578973650932312
➡️  Step 1650 | Geometric Loss: 0.0164 | LR: 1.2e-5 | bv: 7.7574 | dv: 2.4357 | h_norm: 0.9463939666748047
➡️  Step 1700 | Geometric Loss: 0.0166 | LR: 1.2e-5 | bv: 7.7757 | dv: 2.4834 | h_norm: 0.9492393732070923
➡️  Step 1750 | Geometric Loss: 0.016 | LR: 1.2e-5 | bv: 7.7872 | dv: 2.5208 | h_norm: 0.9399847388267517
➡️  Step 1800 | Geometric Loss: 0.0164 | LR: 1.2e-5 | bv: 7.7956 | dv: 2.5447 | h_norm: 0.937726616859436
➡️  Step 1850 | Geometric Loss: 0.0164 | LR: 1.2e-5 | bv: 7.8056 | dv: 2.5674 | h_norm: 0.9349390864372253
➡️  Step 1900 | Geometric Loss: 0.0167 | LR: 1.2e-5 | bv: 7.8106 | dv: 2.5958 | h_norm: 0.9362731575965881
➡️  Step 1950 | Geometric Loss: 0.0167 | LR: 1.2e-5 | bv: 7.829 | dv: 2.6261 | h_norm: 0.9479366540908813
➡️  Step 2000 | Geometric Loss: 0.0163 | LR: 1.2e-5 | bv: 7.8508 | dv: 2.6596 | h_norm: 0.949966311454773
💾 Checkpoint reached! Saving Step 2000...
➡️  Step 2050 | Geometric Loss: 0.0162 | LR: 1.2e-5 | bv: 7.8776 | dv: 2.6897 | h_norm: 0.9537672400474548
➡️  Step 2100 | Geometric Loss: 0.0164 | LR: 1.2e-5 | bv: 7.8921 | dv: 2.7143 | h_norm: 0.9564707279205322