M16: simplify — deduplicate helpers, remove structural copy-paste

- Eliminate deep_cast: express as deep_apply + Nx.as_type closure
- Extract close?/4 and flunk_trajectory/5 to TrainingHelper (was
  copy-pasted across 3 curve tests)
- Extract MNIST load/evaluate helpers to MnistHelper (was duplicated
  between f32 and bf16 test modules)
- Collapse bf16/f16 tensor generators into a single clause that
  delegates to the f32 generator + as_type
- Restore coerce fast path for {:u, 8} (avoids Native.dtype NIF
  call on the most common mismatch type)
- Remove section comments on self-describing private functions
- Single-source LossScaler default scale via @default_scale attribute

Author: ausimian

lib/emily/backend.ex

@@ -55,6 +55,11 @@ defmodule Emily.Backend do
     %{out | data: %B{ref: coerce(ref, type, w)}}
   end
 
+  # Fast path: pred→u8 is the most common mismatch (MLX comparison/logical
+  # ops yield bool; Nx expects u8). The general path covers bf16→f32 type
+  # promotion from Nx.Defn.grad and any other MLX/Nx dtype disagreement.
+  defp coerce(ref, {:u, 8}, w), do: Native.astype(w, ref, {:u, 8})
+
   defp coerce(ref, type, w) do
     case Native.dtype(ref) do
       ^type -> ref

lib/emily/mixed_precision.ex

@@ -62,8 +62,10 @@ defmodule Emily.MixedPrecision do
     after `growth_interval` successful steps it doubles.
     """
 
+    @default_scale 65_536.0
+
     @enforce_keys [:scale]
-    defstruct scale: 65_536.0,
+    defstruct scale: @default_scale,
               growth_factor: 2.0,
               backoff_factor: 0.5,
               growth_interval: 2000,
@@ -81,30 +83,26 @@ defmodule Emily.MixedPrecision do
           :growth_interval,
           :min_scale
         ])
-        |> then(&Keyword.merge(Map.to_list(%__MODULE__{scale: 65_536.0}), &1))
+        |> then(&Keyword.merge(Map.to_list(%__MODULE__{scale: @default_scale}), &1))
 
       struct!(__MODULE__, fields)
     end
   end
 
-  # -------------------------------------------------------------------
-  # Public API
-  # -------------------------------------------------------------------
-
   @doc """
   Downcast float tensors in a nested structure to `type`.
 
   Integer and predicate tensors are left unchanged.
   """
-  def cast_params(params, type), do: deep_cast(params, type)
+  def cast_params(params, type), do: deep_apply(params, &Nx.as_type(&1, type))
 
   @doc """
   Upcast float tensors in a nested gradient structure to `type`.
 
   Semantically identical to `cast_params/2` — exists for readability
   at the call site (the direction of the cast is part of the name).
   """
-  def accumulate_grad(grads, type), do: deep_cast(grads, type)
+  def accumulate_grad(grads, type), do: deep_apply(grads, &Nx.as_type(&1, type))
 
   @doc """
   Create a new dynamic loss scaler.
@@ -167,35 +165,9 @@ defmodule Emily.MixedPrecision do
   """
   def has_overflow?(structure), do: deep_overflow?(structure)
 
-  # -------------------------------------------------------------------
-  # Internal: recursive structure traversal
-  # -------------------------------------------------------------------
-
   defp float_type?({kind, _}) when kind in [:f, :bf], do: true
   defp float_type?(_), do: false
 
-  # deep_cast — cast float tensors to target type
-
-  defp deep_cast(%Nx.Tensor{} = t, type) do
-    if float_type?(t.type), do: Nx.as_type(t, type), else: t
-  end
-
-  defp deep_cast(map, type) when is_map(map) and not is_struct(map) do
-    Map.new(map, fn {k, v} -> {k, deep_cast(v, type)} end)
-  end
-
-  defp deep_cast(tuple, type) when is_tuple(tuple) do
-    tuple |> Tuple.to_list() |> Enum.map(&deep_cast(&1, type)) |> List.to_tuple()
-  end
-
-  defp deep_cast(list, type) when is_list(list) do
-    Enum.map(list, &deep_cast(&1, type))
-  end
-
-  defp deep_cast(other, _type), do: other
-
-  # deep_apply — apply fun to every float tensor
-
   defp deep_apply(%Nx.Tensor{} = t, fun) do
     if float_type?(t.type), do: fun.(t), else: t
   end
@@ -214,8 +186,6 @@ defmodule Emily.MixedPrecision do
 
   defp deep_apply(other, _fun), do: other
 
-  # deep_overflow? — check for nan/inf in any float tensor
-
   defp deep_overflow?(%Nx.Tensor{} = t) do
     if float_type?(t.type) do
       t

test/emily/training/bf16_mlp_curve_test.exs

@@ -13,6 +13,7 @@ defmodule Emily.Training.Bf16MlpCurveTest do
   use ExUnit.Case, async: true
 
   alias Emily.TrainingHelper, as: TH
+  import TH, only: [close?: 4, flunk_trajectory: 5]
 
   @dims {4, 8, 3}
   @batch_shape {16, 4, 3}
@@ -75,18 +76,4 @@ defmodule Emily.Training.Bf16MlpCurveTest do
            "final loss divergence: emily=#{le_final} bin=#{lb_final} " <>
              "reldiff=#{abs(le_final - lb_final) / abs(lb_final)}"
   end
-
-  defp close?(a, b, atol, rtol), do: abs(a - b) <= atol + rtol * abs(b)
-
-  defp flunk_trajectory(i, le, lb, losses_emily, losses_bin) do
-    preview_e = losses_emily |> Enum.take(min(i + 3, length(losses_emily)))
-    preview_b = losses_bin |> Enum.take(min(i + 3, length(losses_bin)))
-
-    flunk("""
-    per-step loss diverged at step #{i}:
-      emily=#{le} bin=#{lb} reldiff=#{abs(le - lb) / abs(lb)}
-    emily trajectory (first #{length(preview_e)} steps): #{inspect(preview_e)}
-    bin   trajectory (first #{length(preview_b)} steps): #{inspect(preview_b)}
-    """)
-  end
 end

test/emily/training/mlp_curve_test.exs

@@ -24,6 +24,7 @@ defmodule Emily.Training.MlpCurveTest do
   use ExUnit.Case, async: true
 
   alias Emily.TrainingHelper, as: TH
+  import TH, only: [close?: 4, flunk_trajectory: 5]
 
   @dims {4, 8, 3}
   @batch_shape {16, 4, 3}
@@ -83,18 +84,4 @@ defmodule Emily.Training.MlpCurveTest do
            "final loss divergence: emily=#{le_final} bin=#{lb_final} " <>
              "reldiff=#{abs(le_final - lb_final) / abs(lb_final)}"
   end
-
-  defp close?(a, b, atol, rtol), do: abs(a - b) <= atol + rtol * abs(b)
-
-  defp flunk_trajectory(i, le, lb, losses_emily, losses_bin) do
-    preview_e = losses_emily |> Enum.take(min(i + 3, length(losses_emily)))
-    preview_b = losses_bin |> Enum.take(min(i + 3, length(losses_bin)))
-
-    flunk("""
-    per-step loss diverged at step #{i}:
-      emily=#{le} bin=#{lb} reldiff=#{abs(le - lb) / abs(lb)}
-    emily trajectory (first #{length(preview_e)} steps): #{inspect(preview_e)}
-    bin   trajectory (first #{length(preview_b)} steps): #{inspect(preview_b)}
-    """)
-  end
 end

test/emily/training/mnist_bf16_full_test.exs

@@ -12,6 +12,8 @@ defmodule Emily.Training.MnistBf16FullTest do
 
   use ExUnit.Case, async: false
 
+  alias Emily.MnistHelper
+
   @moduletag :training_full
   @moduletag capture_log: true
   @moduletag timeout: 600_000
@@ -28,7 +30,7 @@ defmodule Emily.Training.MnistBf16FullTest do
   @target_accuracy 0.955
 
   test "bf16 Axon MLP reaches >#{trunc(@target_accuracy * 100)}% test accuracy under Emily.Compiler" do
-    {train_batches, test_images, test_labels} = load_mnist(@batch_size)
+    {train_batches, test_images, test_labels} = MnistHelper.load_mnist(@batch_size)
 
     policy =
       Axon.MixedPrecision.create_policy(
@@ -51,60 +53,9 @@ defmodule Emily.Training.MnistBf16FullTest do
         compiler: Emily.Compiler
       )
 
-    accuracy = evaluate(model, trained_state, test_images, test_labels)
+    accuracy = MnistHelper.evaluate(model, trained_state, test_images, test_labels)
 
     assert accuracy >= @target_accuracy,
            "bf16 MNIST accuracy #{Float.round(accuracy, 4)} below target #{@target_accuracy}"
   end
-
-  defp load_mnist(batch_size) do
-    {train_images_raw, train_labels_raw} = Scidata.MNIST.download()
-
-    train_images =
-      train_images_raw
-      |> mnist_images_to_tensor()
-      |> Nx.to_batched(batch_size)
-
-    train_labels =
-      train_labels_raw
-      |> mnist_labels_to_tensor()
-      |> Nx.to_batched(batch_size)
-
-    train_batches = Stream.zip(train_images, train_labels)
-
-    {test_images_raw, test_labels_raw} = Scidata.MNIST.download_test()
-
-    test_images = mnist_images_to_tensor(test_images_raw)
-    test_labels = mnist_labels_to_tensor(test_labels_raw)
-
-    {train_batches, test_images, test_labels}
-  end
-
-  defp mnist_images_to_tensor({bin, type, shape}) do
-    bin
-    |> Nx.from_binary(type)
-    |> Nx.reshape(shape)
-    |> Nx.reshape({elem(shape, 0), 784})
-    |> Nx.divide(255.0)
-  end
-
-  defp mnist_labels_to_tensor({bin, type, shape}) do
-    bin
-    |> Nx.from_binary(type)
-    |> Nx.reshape(shape)
-    |> Nx.new_axis(-1)
-    |> Nx.equal(Nx.iota({1, 10}))
-  end
-
-  defp evaluate(model, state, test_images, test_labels) do
-    logits =
-      Axon.predict(model, state, test_images, compiler: Emily.Compiler)
-
-    predicted = Nx.argmax(logits, axis: -1)
-    actual = Nx.argmax(test_labels, axis: -1)
-
-    Nx.mean(Nx.equal(predicted, actual))
-    |> Nx.backend_transfer(Nx.BinaryBackend)
-    |> Nx.to_number()
-  end
 end

test/emily/training/mnist_full_test.exs

@@ -20,6 +20,8 @@ defmodule Emily.Training.MnistFullTest do
 
   use ExUnit.Case, async: false
 
+  alias Emily.MnistHelper
+
   @moduletag :training_full
   @moduletag capture_log: true
   @moduletag timeout: 600_000
@@ -36,7 +38,7 @@ defmodule Emily.Training.MnistFullTest do
   @target_accuracy 0.96
 
   test "Axon MLP reaches >#{trunc(@target_accuracy * 100)}% test accuracy under Emily.Compiler" do
-    {train_batches, test_images, test_labels} = load_mnist(@batch_size)
+    {train_batches, test_images, test_labels} = MnistHelper.load_mnist(@batch_size)
 
     model =
       Axon.input("input", shape: {nil, 784})
@@ -51,67 +53,9 @@ defmodule Emily.Training.MnistFullTest do
         compiler: Emily.Compiler
       )
 
-    accuracy = evaluate(model, trained_state, test_images, test_labels)
+    accuracy = MnistHelper.evaluate(model, trained_state, test_images, test_labels)
 
     assert accuracy >= @target_accuracy,
            "MNIST accuracy #{Float.round(accuracy, 4)} below target #{@target_accuracy}"
   end
-
-  # ---- Data loading ----
-
-  defp load_mnist(batch_size) do
-    # Training set — streamed in batches for Axon.Loop.
-    {train_images_raw, train_labels_raw} = Scidata.MNIST.download()
-
-    train_images =
-      train_images_raw
-      |> mnist_images_to_tensor()
-      |> Nx.to_batched(batch_size)
-
-    train_labels =
-      train_labels_raw
-      |> mnist_labels_to_tensor()
-      |> Nx.to_batched(batch_size)
-
-    train_batches = Stream.zip(train_images, train_labels)
-
-    # Test set — loaded as whole tensors for one-shot evaluation.
-    {test_images_raw, test_labels_raw} = Scidata.MNIST.download_test()
-
-    test_images = mnist_images_to_tensor(test_images_raw)
-    test_labels = mnist_labels_to_tensor(test_labels_raw)
-
-    {train_batches, test_images, test_labels}
-  end
-
-  defp mnist_images_to_tensor({bin, type, shape}) do
-    bin
-    |> Nx.from_binary(type)
-    |> Nx.reshape(shape)
-    # {N, 1, 28, 28} → {N, 784} + normalize to [0, 1].
-    |> Nx.reshape({elem(shape, 0), 784})
-    |> Nx.divide(255.0)
-  end
-
-  defp mnist_labels_to_tensor({bin, type, shape}) do
-    bin
-    |> Nx.from_binary(type)
-    |> Nx.reshape(shape)
-    |> Nx.new_axis(-1)
-    |> Nx.equal(Nx.iota({1, 10}))
-  end
-
-  # ---- Evaluation ----
-
-  defp evaluate(model, state, test_images, test_labels) do
-    logits =
-      Axon.predict(model, state, test_images, compiler: Emily.Compiler)
-
-    predicted = Nx.argmax(logits, axis: -1)
-    actual = Nx.argmax(test_labels, axis: -1)
-
-    Nx.mean(Nx.equal(predicted, actual))
-    |> Nx.backend_transfer(Nx.BinaryBackend)
-    |> Nx.to_number()
-  end
 end

test/emily/training/transformer_block_curve_test.exs

@@ -16,6 +16,7 @@ defmodule Emily.Training.TransformerBlockCurveTest do
   use ExUnit.Case, async: true
 
   alias Emily.TrainingHelper, as: TH
+  import TH, only: [close?: 4, flunk_trajectory: 5]
 
   @embed 16
   @ff 32
@@ -75,18 +76,4 @@ defmodule Emily.Training.TransformerBlockCurveTest do
            "final loss divergence: emily=#{le_final} bin=#{lb_final} " <>
              "reldiff=#{abs(le_final - lb_final) / abs(lb_final)}"
   end
-
-  defp close?(a, b, atol, rtol), do: abs(a - b) <= atol + rtol * abs(b)
-
-  defp flunk_trajectory(i, le, lb, losses_emily, losses_bin) do
-    preview_e = losses_emily |> Enum.take(min(i + 3, length(losses_emily)))
-    preview_b = losses_bin |> Enum.take(min(i + 3, length(losses_bin)))
-
-    flunk("""
-    per-step transformer-block loss diverged at step #{i}:
-      emily=#{le} bin=#{lb} reldiff=#{abs(le - lb) / abs(lb)}
-    emily trajectory (first #{length(preview_e)} steps): #{inspect(preview_e)}
-    bin   trajectory (first #{length(preview_b)} steps): #{inspect(preview_b)}
-    """)
-  end
 end

test/support/backend_generators.ex

@@ -40,18 +40,8 @@ defmodule Emily.BackendGenerators do
     |> map(&Nx.reshape(&1, shape))
   end
 
-  def tensor(shape, {:bf, 16}) do
-    list_of(float(min: -10.0, max: 10.0), length: Nx.size(shape))
-    |> map(&Nx.tensor(&1, type: {:f, 32}, backend: Nx.BinaryBackend))
-    |> map(&Nx.reshape(&1, shape))
-    |> map(&Nx.as_type(&1, {:bf, 16}))
-  end
-
-  def tensor(shape, {:f, 16}) do
-    list_of(float(min: -10.0, max: 10.0), length: Nx.size(shape))
-    |> map(&Nx.tensor(&1, type: {:f, 32}, backend: Nx.BinaryBackend))
-    |> map(&Nx.reshape(&1, shape))
-    |> map(&Nx.as_type(&1, {:f, 16}))
+  def tensor(shape, {kind, 16} = type) when kind in [:bf, :f] do
+    tensor(shape, {:f, 32}) |> map(&Nx.as_type(&1, type))
   end
 
   def tensor(shape, {:s, bits}) do