E2A2 — stereo audio

E2A2 (E²A², Encoding Efficient Asymmetric Autoencoders, github.com/danjacobellis/E2A2) is the successor of FRAPPE v3: the same split trainer (frozen per-group encoders — one learned strided projection per latent channel — SC6 softsign companding onto the coder alphabet [−31, 31], one merged decoder per operating point) with the entropy stage replaced by GGDLPC (a learned FIR predictor per channel, two affine laws mapping a causal statistic to a folded-GGD Huffman table, plus a derived zero-run mode; the C engine builds itself on import). Every reported rate is real byte-padded stream bits, round-trip verified in this notebook by decoding the latents back from the truncated bytes.

E2A2A is the E2A2 stereo-music codec: a 7-group rate ladder (ps 256/256/128/64/32/16/8, group sizes 4/5/6/6/12/9/9, macroregion 16384 samples ≈ 372 ms at 44.1 kHz), trained on danjacobellis/musdb_segments.

Stream layout all_detail: no global stream — every channel is coded per macroregion, and every macroregion is decodable from its own bytes alone (one byte-pad each). The stream for operating point n is the literal bit-prefix covering channels ≤ n.

The codec is loaded from the Hugging Face hub (danjacobellis/E2A2, audio/ subdirectory) through the compressors.e2a2 public API — there is no local checkpoint or training code in this notebook.

import json, glob
import torch, numpy as np, matplotlib.pyplot as plt
import torchaudio
from IPython.display import Audio, display
from datasets import load_dataset
from torchvision.transforms import ToPILImage
from huggingface_hub import hf_hub_download
from compressors import e2a2

Load the codec and its entropy coder

One MergedAutoencoder per operating point, all from the hub. Each operating point has its own merged decoder (multidec); the encoders froze at merge, so every operating point’s encoder is a channel prefix of the full one. entropy.json is the GGDLPC blob (per-channel predictor taps + law scalars); the Huffman tables are generated from it when the coder loads.

NAME = 'audio'
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
config, models, coder = e2a2.load_all_codecs(NAME, device=device)
ops = list(config.cumulative_channels)
model = models[ops[-1]]
rom = e2a2.rom_facts(coder.blob)

print(f'codec         = {config.codec}/{NAME}   (hub format {config.format}, source checkpoint {config.source_checkpoint})')
print(f'modality      = {config.modality}   spatial rank = {config.dim}   input channels = {config.input_channels}')
print(f'ps ladder     = {list(config.ps)} x {list(config.group_sizes)} channels')
print(f'boundaries    = {ops}  (cumulative channels = operating points)')
for n in ops:
    d = config.decoder_configs[ops.index(n)]
    print(f'  op {n:>2}: decoder ps={d["decoder_ps"]} dim={d["decoder_dim"]} depth={d["decoder_depth"]}')
for s, (ps_s, start, end) in enumerate(model.scale_groups):
    print(f'  scale {s+1}: ps={ps_s:>4}  channels {start+1}-{end}')
print(f'\nentropy coder: GGDLPC, stream_mode={coder.stream_mode}, macroregion={coder.macroregion} samples'
      + (f', global stream = channels 1-{coder.n_coarse}' if coder.n_coarse else ', no global stream'))
print(f'  alphabet [-31, 31]; persistent model = {rom["law_scalars"]} law scalars + '
      f'{rom["predictor_ints"]} predictor ints for {rom["n_channels"]} channels')
print(f'  {rom["note"]}')

codec         = E2A2/audio   (hub format e2a2_hub1, source checkpoint e2a2a_g7.pth)
modality      = audio   spatial rank = 1   input channels = 2
ps ladder     = [256, 256, 128, 64, 32, 16, 8] x [4, 5, 6, 6, 12, 9, 9] channels
boundaries    = [4, 9, 15, 21, 33, 42, 51]  (cumulative channels = operating points)
  op  4: decoder ps=128 dim=1024 depth=12
  op  9: decoder ps=128 dim=1024 depth=12
  op 15: decoder ps=128 dim=1024 depth=24
  op 21: decoder ps=128 dim=1024 depth=24
  op 33: decoder ps=128 dim=1024 depth=24
  op 42: decoder ps=128 dim=1024 depth=24
  op 51: decoder ps=128 dim=1024 depth=24
  scale 1: ps= 256  channels 1-9
  scale 2: ps= 128  channels 10-15
  scale 3: ps=  64  channels 16-21
  scale 4: ps=  32  channels 22-33
  scale 5: ps=  16  channels 34-42
  scale 6: ps=   8  channels 43-51

entropy coder: GGDLPC, stream_mode=all_detail, macroregion=16384 samples, no global stream
  alphabet [-31, 31]; persistent model = 306 law scalars + 153 predictor ints for 51 channels
  edges, folded GGD pmfs, quotient split, canonical Huffman tables and the run mode are GENERATED from the per-channel scalars at load (GGDLPC.tables)

Analysis filterbank

One learned (2, ps) projection per latent channel: impulse response (left, averaged over the two input channels) and magnitude response (right); ≤ 6 channels per panel.

def chunk_sizes(n, max_per=6):
    k = max(1, -(-n // max_per))
    return [n // k + (1 if i < n % k else 0) for i in range(k)]

fs_hz = config.sample_rate
N = 3001
freqs = np.fft.fftshift(np.fft.fftfreq(N)) * fs_hz

panels = []
for s, (ps_s, start, end) in enumerate(model.scale_groups):
    W = model.encoders[s][0].weight.data.cpu().numpy().mean(1)
    off = 0
    for sz in chunk_sizes(end - start):
        panels.append((ps_s, start + off + 1, start + off + sz, W[off:off + sz]))
        off += sz

fig, axes = plt.subplots(len(panels), 2, figsize=(10, 1.7 * len(panels)), dpi=100, constrained_layout=True)
axes = np.atleast_2d(axes)
for (ax_imp, ax_mag), (ps_s, ch_lo, ch_hi, filt) in zip(axes, panels):
    t_ms = np.arange(filt.shape[1]) / fs_hz * 1e3
    for c, h in enumerate(filt, start=ch_lo):
        H = np.fft.fftshift(np.fft.fft(h, N))
        ax_imp.plot(t_ms, h, lw=1.0, label=f'ch{c}')
        ax_mag.plot(freqs, np.abs(H), lw=1.0, label=f'ch{c}')
    ax_imp.set_title(f'ps={ps_s}   channels {ch_lo}-{ch_hi}  (impulse response)', fontsize=9)
    ax_mag.set_title(f'ps={ps_s}   channels {ch_lo}-{ch_hi}', fontsize=9)
    ax_mag.set_xlim(-fs_hz / 2, fs_hz / 2)
    ax_imp.set_ylabel(r'$h(t)$'); ax_mag.set_ylabel(r'$|H(f)|$')
    ax_mag.legend(fontsize=7, loc='upper right', framealpha=0.6, ncol=2)
axes[-1, 0].set_xlabel('time (ms)'); axes[-1, 1].set_xlabel('frequency (Hz)')
fig.suptitle('analysis filterbank', fontsize=11)
plt.show()

Load an example clip

The danjacobellis/music dataset (private due to copyright) holds 24 stereo music tracks (lossy MP3 sources at near-lossless rates). A five-second stereo 44.1 kHz excerpt, zero-mean peak-normalized to [-1, 1] (the trainer’s validation convention), cropped to a whole number of macroregions.

CLIP_INDEX = 16
START_SECONDS = 150.0
DURATION_SECONDS = 5.0

ds = load_dataset("danjacobellis/music", split="train")
samples = ds[CLIP_INDEX]["audio"].get_all_samples()
fs = samples.sample_rate
assert fs == config.sample_rate, (fs, config.sample_rate)
i0 = round(START_SECONDS * fs)
x = samples.data[:, i0 : i0 + round(DURATION_SECONDS * fs)].to(torch.float)

x = x - x.mean()
x = x / (x.abs().max() + 1e-8)
unit = max(coder.macroregion, max(config.ps))
L = unit * (x.shape[-1] // unit)
x = x[..., :L]
n_samples = x.shape[-1]
x_in = x.unsqueeze(0).to(device)
print(f'fs       = {fs} Hz')
print(f'clip     = {tuple(x.shape)}  ({n_samples / fs:.1f} s),  range [{x.min():.2f}, {x.max():.2f}]')
print(f'model in = {tuple(x_in.shape)}  ({n_samples // coder.macroregion} macroregions of '
      f'{coder.macroregion} samples = {coder.macroregion / fs * 1e3:.0f} ms each)')

fs       = 44100 Hz
clip     = (2, 212992)  (4.8 s),  range [-1.00, 0.95]
model in = (1, 2, 212992)  (13 macroregions of 16384 samples = 372 ms each)

mel = torchaudio.transforms.MelSpectrogram(sample_rate=fs, n_fft=4096, hop_length=512)

def make_spectrogram(x_1ch):
    S = mel(x_1ch).log()
    S = ((S - S.mean()) / (3 * S.std()) + 0.5).clamp(0, 1)
    return ToPILImage()(S.flip(0))

for ch, name in enumerate(['left', 'right']):
    print(f'{name} channel')
    display(make_spectrogram(x[ch]))

Audio(x.numpy(), rate=fs)

left channel

right channel

Analysis transform and rounding

Encode to the per-scale latents (companding included) and round to the coder’s [-31, 31] grid.

with torch.no_grad():
    latents = model.encode(x_in)
n_latent_values = sum(z.numel() for z in latents)
print(f'{x_in.numel()} input samples -> {n_latent_values} latent values ({n_latent_values / x_in.numel():.3f}x)')
for s, z in enumerate(latents):
    print(f'  scale {s+1} (ps={model.scale_groups[s][0]:>4}): latent {tuple(z.shape)}  range [{z.min():.1f}, {z.max():.1f}]')
n_clip = sum(int((z.round().abs() > 31).sum()) for z in latents)
latents_q = e2a2.encode_to_latents(model, x_in)
latents_np = [z[0].numpy() for z in latents_q]
print(f'clipped: {n_clip}/{n_latent_values}')

plt.figure(figsize=(5, 2), dpi=120)
plt.hist(np.concatenate([z.ravel() for z in latents_np]), range=(-31.5, 31.5), bins=63, width=0.85)
plt.xlim([-15, 15]); plt.title('histogram of integer latents (zoomed)')
plt.xlabel('value'); plt.ylabel('count'); plt.tight_layout(); plt.show()

425984 input samples -> 476736 latent values (1.119x)
  scale 1 (ps= 256): latent (1, 9, 832)  range [-5.3, 5.6]
  scale 2 (ps= 128): latent (1, 6, 1664)  range [-4.6, 5.4]
  scale 3 (ps=  64): latent (1, 6, 3328)  range [-11.4, 11.6]
  scale 4 (ps=  32): latent (1, 12, 6656)  range [-11.2, 11.7]
  scale 5 (ps=  16): latent (1, 9, 13312)  range [-2.3, 2.2]
  scale 6 (ps=   8): latent (1, 9, 26624)  range [-17.1, 16.4]
clipped: 0/476736

GGDLPC all-detail entropy coding

encode_latents produces one stream per macroregion and nothing else — every macroregion is independently decodable from its own bytes alone. The stream for operating point n is the literal bit-prefix covering channels ≤ n. CR is against 16-bit stereo PCM (32 bits per frame).

enc = e2a2.encode_latents(coder, latents_q)
assert enc.gblob == b''
region_bytes = (enc.prefix[:, -1] + 7) // 8
print(f'{len(enc.streams)} macroregion streams, one byte-pad each; at max detail '
      f'{int(region_bytes.sum()):,} bytes total, per region min/median/max = '
      f'{int(region_bytes.min())}/{int(np.median(region_bytes))}/{int(region_bytes.max())}')
print(f'\n{"op":>4} {"stream bits":>12} {"bits/frame":>11} {"kbps":>8} {"CR":>8} {"bits/latent":>12}')
for n in ops:
    bits = e2a2.op_bits(coder, enc, n)
    n_lat_n = sum(min(max(n - start, 0), end - start) * z.shape[1]
                  for z, (ps_s, start, end) in zip(latents_np, model.scale_groups))
    kbps_n = bits / (n_samples / fs) / 1000
    print(f'{n:>4} {bits:>12,} {bits / n_samples:>11.4f} {kbps_n:>8.2f} {32 * n_samples / bits:>7.1f}x {bits / n_lat_n:>12.3f}')

13 macroregion streams, one byte-pad each; at max detail 117,198 bytes total, per region min/median/max = 6771/9672/10724

  op  stream bits  bits/frame     kbps       CR  bits/latent
   4        8,408      0.0395     1.74   810.6x        2.526
   9       18,072      0.0848     3.74   377.1x        2.413
  15       40,656      0.1909     8.42   167.6x        2.327
  21      102,920      0.4832    21.31    66.2x        2.749
  33      330,832      1.5533    68.50    20.6x        2.820
  42      422,624      1.9842    87.50    16.1x        1.782
  51      937,584      4.4020   194.13     7.3x        1.967

Round-trip verification

verify_latents decodes the latents back from the truncated bytes alone at every operating point’s literal bit-prefix (the global stream from its own bytes; every macroregion from its own bytes, no cross-region input) and asserts them bit-exact against the encoder’s. The reconstructions below are decoded from those byte-decoded latents (decode_latents), not from the encoder’s copy.

e2a2.verify_latents(coder, enc, latents_q, ops)
print(f'round trip verified at ops {ops}: every op decoded from its truncated bytes alone, bit-exact')

round trip verified at ops [4, 9, 15, 21, 33, 42, 51]: every op decoded from its truncated bytes alone, bit-exact

Reconstructions at each operating point

One encode serves every operating point: the latents for op n are decoded from the truncated bytes and fed to that point’s own merged decoder; the rate is the same encode’s literal per-region bit-prefix. Per point: the left channel’s log-mel spectrogram and the stereo player. PSNR is on the [0, 1]-mapped signals (x/2 + 0.5).

print('original (left channel)')
display(make_spectrogram(x[0]))
rd_clip = []
for n in ops:
    lts = e2a2.decode_latents(coder, enc, n)
    with torch.no_grad():
        xh = models[n].decode([z.to(device) for z in lts]).clamp(-1, 1)[0].float().cpu()
    bits_n = e2a2.op_bits(coder, enc, n)
    kbps_n = bits_n / (n_samples / fs) / 1000
    psnr_n = -10 * torch.nn.functional.mse_loss(x / 2 + 0.5, xh / 2 + 0.5).log10().item()
    rd_clip.append((bits_n / n_samples, psnr_n, n))
    print(f'{n:>2} ch:  {bits_n / n_samples:.4f} bits/frame   {kbps_n:8.2f} kbps   '
          f'CR={32 * n_samples / bits_n:8.1f}x   PSNR={psnr_n:.2f} dB')
    display(make_spectrogram(xh[0]))
    display(Audio(xh.numpy(), rate=fs))

original (left channel)

 4 ch:  0.0395 bits/frame       1.74 kbps   CR=   810.6x   PSNR=23.12 dB

 9 ch:  0.0848 bits/frame       3.74 kbps   CR=   377.1x   PSNR=25.28 dB

15 ch:  0.1909 bits/frame       8.42 kbps   CR=   167.6x   PSNR=28.80 dB

21 ch:  0.4832 bits/frame      21.31 kbps   CR=    66.2x   PSNR=32.69 dB

33 ch:  1.5533 bits/frame      68.50 kbps   CR=    20.6x   PSNR=42.37 dB

42 ch:  1.9842 bits/frame      87.50 kbps   CR=    16.1x   PSNR=42.73 dB

51 ch:  4.4020 bits/frame     194.13 kbps   CR=     7.3x   PSNR=47.39 dB

Per-macroregion rate profile

Each macroregion carries its own byte-padded stream — the rate allocation over time is directly visible, and any region can be transmitted, dropped, or upgraded independently.

t_edges = np.arange(len(enc.streams) + 1) * coder.macroregion / fs
fig, axes = plt.subplots(2, 1, figsize=(8, 4), dpi=120, sharex=True, gridspec_kw={'height_ratios': [1, 1.6]})
t_sig = np.arange(n_samples) / fs
axes[0].plot(t_sig, x[0].numpy(), lw=0.3, color='0.4')
axes[0].set_ylabel('amplitude'); axes[0].set_title('input (left channel)', fontsize=9)
for n in ops:
    rb = 8 * ((enc.prefix[:, n - 1] + 7) // 8)
    axes[1].stairs(rb / (coder.macroregion / fs) / 1000, t_edges, label=f'{n} ch')
axes[1].set_xlabel('time (s)'); axes[1].set_ylabel('kbps')
axes[1].set_title(f'per-macroregion rate ({coder.macroregion} samples each)', fontsize=9)
axes[1].legend(fontsize=8, ncol=2); axes[1].grid(True, alpha=0.3)
plt.tight_layout(); plt.show()

Rate–distortion

This clip’s operating points over the whole-musdb-validation curve (262 clips) from results/e2a2_audio/ (falls back to the hub rd.json), against FRAPPE v3 audio (hub rd.json) and the FRAPPE v2 audio baseline (run B1) from results/frappe_v2_audio/. Rates in bits per frame (both channels).

def latest(pattern):
    paths = sorted(glob.glob(pattern))
    return paths[-1] if paths else None

rd_path = latest(f'results/e2a2_{NAME}/rate_distortion_*.json')
if rd_path is not None:
    rd = json.load(open(rd_path))
    key = 'channel_counts' if 'channel_counts' in rd else 'quality_values'
    rd_ops = list(rd[key])
    rd_label = f"E2A2 {NAME} — whole validation split ({rd.get('n_images', rd.get('n_samples'))} samples, compressors harness, verified streams)"
else:
    rd = None
    hub_rd = e2a2.hub_rd_curve(NAME)
    rd_ops = [p['n_ch'] for p in hub_rd['rd_curve']]
    rd_label = f"E2A2 {NAME} — whole validation split (hub rd.json, trainer's verified streams)"

if rd is not None:
    val_bpf = [rd['results'][str(n)]['mean']['kbps'] * 1000 / rd['sample_rate'] for n in rd_ops]
    val_psnr = [rd['results'][str(n)]['mean']['PSNR_dB'] for n in rd_ops]
else:
    val_bpf = [p['bpp'] for p in hub_rd['rd_curve']]; val_psnr = [p['psnr'] for p in hub_rd['rd_curve']]

rd3 = json.load(open(hf_hub_download(repo_id='danjacobellis/FRAPPEv3', filename='audio/rd.json')))
v3_bpf = [p['bpp'] for p in rd3['rd_curve']]; v3_psnr = [p['psnr'] for p in rd3['rd_curve']]
rd2 = json.load(open(latest('results/frappe_v2_audio/rate_distortion_*.json')))
v2_bpf = [rd2['results'][str(q)]['mean']['kbps'] * 1000 / rd2['sample_rate'] for q in rd2['quality_values']]
v2_psnr = [rd2['results'][str(q)]['mean']['PSNR_dB'] for q in rd2['quality_values']]

plt.figure(figsize=(6.5, 4), dpi=120)
plt.semilogx(v2_bpf, v2_psnr, '^-', color='black', ms=5, lw=1.2, alpha=0.8, label='FRAPPE v2 audio + JPEG-LS (B1)')
plt.semilogx(v3_bpf, v3_psnr, 'D-', color='tab:gray', ms=4, lw=1.0, label='FRAPPE v3 audio (lawcoder, hub rd.json)')
plt.semilogx(val_bpf, val_psnr, 'o-', color='tab:blue', label=rd_label)
plt.semilogx([p[0] for p in rd_clip], [p[1] for p in rd_clip], 's--', color='tab:blue', alpha=0.5, label=f'E2A2 audio — this clip (#{CLIP_INDEX})')
for n, bv, psv in zip(rd_ops, val_bpf, val_psnr):
    plt.annotate(f'{n}ch', (bv, psv), fontsize=7, color='tab:blue', xytext=(3, -9), textcoords='offset points')
plt.xlabel('bits per frame (2 channels)'); plt.ylabel('PSNR (dB)')
plt.title('E2A2 audio vs FRAPPE v2/v3 — musdb validation')
plt.grid(True, which='both', alpha=0.3); plt.legend(fontsize=7); plt.tight_layout(); plt.show()

print('E2A2 audio, musdb validation average (verified streams):')
for n, bv, psv in zip(rd_ops, val_bpf, val_psnr):
    print(f'  {n:>2} ch:  {bv:.6f} bits/frame  ({bv * fs / 1000:8.3f} kbps)   PSNR={psv:.3f} dB')

E2A2 audio, musdb validation average (verified streams):
   4 ch:  0.032025 bits/frame  (   1.412 kbps)   PSNR=28.061 dB
   9 ch:  0.056569 bits/frame  (   2.495 kbps)   PSNR=29.815 dB
  15 ch:  0.106056 bits/frame  (   4.677 kbps)   PSNR=31.242 dB
  21 ch:  0.319616 bits/frame  (  14.095 kbps)   PSNR=33.885 dB
  33 ch:  1.085801 bits/frame  (  47.884 kbps)   PSNR=39.369 dB
  42 ch:  1.758319 bits/frame  (  77.542 kbps)   PSNR=41.865 dB
  51 ch:  3.603938 bits/frame  ( 158.934 kbps)   PSNR=45.764 dB