desink.livemd

Diffusion Encoder - sinkhorn

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

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

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

Diffusion

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.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.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 \\ 30) do
    # np = DLM.Loss.l2_normalize(pred)
    # nt = DLM.Loss.l2_normalize(target)
    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.ContextualEmbedTrainer 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. The Sinkhorn Loss Function
  defn compute_loss(e, batch_tokens) do
    # Extract inputs (t) and targets (t+1) from the sequence
    seq_len = Nx.axis_size(batch_tokens, 1) - 1
    inputs = Nx.slice_along_axis(batch_tokens, 0, seq_len, axis: 1)
    targets = Nx.slice_along_axis(batch_tokens, 1, seq_len, axis: 1)

    # Map integers to their current 512D coordinates
    pred_coords = Nx.take(e, inputs)
    target_coords = Nx.take(e, targets)

    # Flatten the sequence and batch dimensions together into a single cloud of points
    flat_preds = Nx.reshape(pred_coords, {:auto, 512})
    flat_targets = Nx.reshape(target_coords, {:auto, 512})
    
    # Normalize to ensure everything remains perfectly on the surface of the void
    p_norm = l2_normalize(flat_preds)
    t_norm = l2_normalize(flat_targets)

    # Execute Optimal Transport (epsilon = 0.1 to match your void's scale)
    DLM.SinkhornOT.compute(p_norm, t_norm, 0.1)
  end

  # ⚡ 2. Newton-Schulz Iteration for Orthogonal Isometry
  defn newton_schulz(g) do
    norm = Nx.sqrt(Nx.sum(Nx.pow(g, 2)))
    x_init = Nx.divide(g, Nx.add(norm, 1.0e-8))

    {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

  # ⚡ 3. The Muon Update Step
  defn muon_step(p, g, m, lr, max_norm \\ 20.0) do
    beta1 = 0.95
    new_m = beta1 * m + (1.0 - beta1) * g
    ortho_update = newton_schulz(new_m)
    
    r = Nx.axis_size(p, 0)
    c = Nx.axis_size(p, 1)
    scale = Nx.divide(Nx.max(r, c), 5.0) |> Nx.as_type(:f32)
    
    p_new = p - lr * scale * ortho_update
    
    # Clamp norm to prevent runaway expansion
    current_norm = Nx.sqrt(Nx.sum(Nx.pow(p_new, 2)))
    p_clipped = Nx.select(
      current_norm > max_norm,
      p_new * (max_norm / (current_norm + 1.0e-8)),
      p_new
    )
    
    {p_clipped, new_m}
  end

  # ⚡ 4. The Unified Master Step
  defn compute_grad_and_step(e, m, batch_tokens, _step_t, lr) do
    # Calculate the global transport cost and the exact gradient to minimize it
    {loss, grad_e} = value_and_grad(e, fn e_params -> 
      compute_loss(e_params, batch_tokens) 
    end)

    # Pass the gradient through the Muon forge to strictly enforce orthogonality
    {new_e, new_m} = DLM.SparseSinkGD.step(e, grad_e, m, lr)

    {loss, new_e, new_m}
  end
end

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

batch_size = 32
seq_len = 128
key = Nx.Random.key(42) |> Nx.backend_copy(gpu)

embeds = Nx.Random.normal(key, 0.0, 0.1, shape: {128, 512}) |> elem(0) |> Nx.backend_copy(gpu)
m_embeds = Nx.broadcast(0.0, {128, 512}) |> Nx.backend_copy(gpu)
lr = Nx.tensor(3.0e-5, type: :f32) |> Nx.backend_copy(gpu)

# Flatten the 2D dataset to 1D for the streamer
full_dataset_2d = File.read!("scholar_full_mixed_128.bin") |> :erlang.binary_to_term()
dataset_1d = Nx.flatten(full_dataset_2d)
data_stream = DLM.DataStreamer.build_infinite_stream(dataset_1d, batch_size, seq_len)

total_steps = 13206
IO.puts("🚀 IGNITION: Phase 1 Contextual Shape Training...")

{final_embeds, _final_m} =
  Enum.reduce(
    Enum.zip(0..(total_steps - 1), data_stream),
    {embeds, m_embeds},
    fn {step, batch}, {e, m} ->
      gpu_batch = Nx.backend_copy(batch, gpu) |> Nx.as_type(:s32)
      step_t = Nx.tensor(step, type: :f32) |> Nx.backend_copy(gpu)

      {loss, new_e, new_m} =
        DLM.ContextualEmbedTrainer.compute_grad_and_step(e, m, gpu_batch, step_t, lr)

      if rem(step, 100) == 0 do
        IO.puts("Step #{step} | Loss: #{Nx.to_number(loss)}")
      end

      {new_e, new_m}
    end
  )

IO.puts("🎉 PHASE 1 TRAINING COMPLETE!")

# Save — embeddings only, matching the format your downstream code expects
cpu_embeds = Nx.backend_copy(final_embeds, Nx.BinaryBackend)

# ⚡ Match the tuple format your Phase 2 loader expects:
# %{params: phase1_p} -> {cpu_embeds, _w1, _w2, _key}
dummy = Nx.broadcast(0.0, {1}) |> Nx.backend_copy(Nx.BinaryBackend)
File.write!(
  "phase1_sink_diffusion_13k_COMPLETE.bin",
  :erlang.term_to_binary(%{params: {cpu_embeds, dummy, dummy, dummy}})
)

22:16:57.349 [info] XLA service 0x7619280b6080 initialized for platform ROCM (this does not guarantee that XLA will be used). Devices:

22:16:57.351 [info]   StreamExecutor [0]: AMD Radeon RX 7600, AMDGPU ISA version: gfx1100 (Driver: 6.4.43482; Runtime: 6.4.43482; Toolkit: 6.4.43482; DNN: 0.0.0)

22:16:57.351 [info] Using BFC allocator.

22:16:57.351 [info] XLA backend will use up to 6429868032 bytes on device 0 for BFCAllocator.

22:16:57.351 [info] XLA backend will use up to 2143289344 bytes on device 0 for CollectiveBFCAllocator.
🚀 IGNITION: Phase 1 Contextual Shape Training...

22:16:57.913 [info] Merging Dots in computation: region_17.18
Step 0 | Loss: 4.953514007866033e-7
Step 100 | Loss: 4.928349994770542e-7
Step 200 | Loss: 1.0723611012508627e-5
Step 300 | Loss: 1.3578316429629922e-4
Step 400 | Loss: 1.2233336747158319e-5

# =====================================================================
# Phase 1: Manifold Visualization
# =====================================================================
alias VegaLite, as: Vl

# 1. Extract the learned 128x512 embeddings from your final params tuple
# {final_embeds, _w1, _w2, _key} = final_p

# Bring it back to the CPU for LinAlg processing
coords = Nx.backend_copy(final_embeds, Nx.BinaryBackend)

# 2. PCA: Squash 512D down to 2D
mean = Nx.mean(coords, axes: [0], keep_axes: true)
centered = Nx.subtract(coords, mean)

# Perform Singular Value Decomposition
{_u, _s, vt} = Nx.LinAlg.svd(centered)

# Extract the top 2 Principal Components
v_top2 = Nx.transpose(vt)[[.., 0..1]]

# Project the 512D coordinates onto the new 2D plane
coords_2d = Nx.dot(centered, v_top2)

# 3. Extract X and Y for graphing
x_vals = coords_2d[[.., 0]] |> Nx.to_flat_list()
y_vals = coords_2d[[.., 1]] |> Nx.to_flat_list()

# 4. Tag and categorize the printable ASCII characters
plot_data =
  Enum.map(32..126, fn ascii ->
    char = List.to_string([ascii])
    
    type = cond do
      char == " " -> "Space"
      char =~ ~r/[aeiouAEIOU]/ -> "Vowel"
      char =~ ~r/[a-zA-Z]/ -> "Consonant"
      char =~ ~r/[0-9]/ -> "Number"
      true -> "Punctuation"
    end

    %{
      "character" => char,
      "x" => Enum.at(x_vals, ascii),
      "y" => Enum.at(y_vals, ascii),
      "type" => type
    }
  end)

# 5. Render the Latent Map
Vl.new(width: 800, height: 600, title: "Diffusion-Molded 512D Latent Manifold")
|> Vl.data_from_values(plot_data)
|> Vl.mark(:text, size: 16, font: "monospace", font_weight: "bold")
|> Vl.encode_field(:x, "x", type: :quantitative, title: "Principal Component 1")
|> Vl.encode_field(:y, "y", type: :quantitative, title: "Principal Component 2")
|> Vl.encode_field(:text, "character", type: :nominal)
|> Vl.encode_field(:color, "type", type: :nominal, scale: [
  domain: ["Space", "Vowel", "Consonant", "Number", "Punctuation"],
  range: ["#FF0000", "#00AEEF", "#2A363B", "#99B898", "#E84A5F"]
])
|> Kino.VegaLite.new()

17:29:25.396 [info] Merging Dots in computation: region_7.7

17:29:25.396 [info] Merging Dots in computation: region_16.14

17:29:25.396 [info] Merging Dots in computation: region_54.58

17:29:25.396 [info] Merging Dots in computation: region_6.61

Probe

defmodule DLM.LatentRadar do
  import Nx.Defn

  # Calculate the Squared Euclidean distance between all points simultaneously
  defn pairwise_squared_distances(coords) do
    a = Nx.new_axis(coords, 1)
    b = Nx.new_axis(coords, 0)
    
    diff = Nx.subtract(a, b)
    # ⚡ No Nx.sqrt() here! We match the Sinkhorn cost_matrix exactly.
    Nx.sum(Nx.pow(diff, 2), axes: [-1]) 
  end
end

# =====================================================================
# Execution & Analysis
# =====================================================================

gpu = {EXLA.Backend, client: :rocm}

# 1. Load the frozen dictionary directly
%{params: {cpu_embeds, _, _, _}} = File.read!("phase1_msenorm_diffusion_25k_COMPLETE.bin") |> :erlang.binary_to_term()

# 2. L2 Normalize it (This is exactly how Phase 2 sees the targets)
frozen_map = cpu_embeds |> Nx.backend_copy(gpu) |> DLM.ContextualEmbedTrainer.l2_normalize()

vocab_size = Nx.axis_size(frozen_map, 0)

# 3. Generate the 256x256 Distance Matrix
dist_matrix = DLM.LatentRadar.pairwise_squared_distances(frozen_map)

# 4. Mask the diagonal (distance from a character to itself is 0.0)
mask = Nx.eye({vocab_size, vocab_size}) |> Nx.multiply(9999.0) |> Nx.backend_copy(gpu)
masked_dists = Nx.add(dist_matrix, mask)

# Extract the vital statistics
min_dist = Nx.reduce_min(masked_dists) |> Nx.to_number()
max_dist = Nx.reduce_max(dist_matrix) |> Nx.to_number()
avg_dist = Nx.mean(dist_matrix) |> Nx.to_number() 

IO.puts "🌌 512D Latent Void Geometry (Squared Euclidean) 🌌"
IO.puts "=================================================="
IO.puts "Absolute Closest Neighbors: #{Float.round(min_dist, 4)}"
IO.puts "Furthest Two Characters:    #{Float.round(max_dist, 4)}"
IO.puts "Average Distance:           #{Float.round(avg_dist, 4)}\n"

# Let's inspect the letter 'e' (ASCII 101)
e_ascii = 101
e_masked = masked_dists[e_ascii]
closest_idx = Nx.argmin(e_masked) |> Nx.to_number()
closest_dist = e_masked[closest_idx] |> Nx.to_number()

IO.puts "🔍 Probe: The Letter 'e' (ASCII #{e_ascii})"
IO.puts "Closest neighbor is ASCII #{closest_idx} (Squared Dist: #{Float.round(closest_dist, 4)})"

🌌 512D Latent Void Geometry (Squared Euclidean) 🌌
==================================================
Absolute Closest Neighbors: 0.0
Furthest Two Characters:    2.2859
Average Distance:           0.9115

🔍 Probe: The Letter 'e' (ASCII 101)
Closest neighbor is ASCII 52 (Squared Dist: 0.0)

:ok