---
format: typebulb/v1
name: NeRF
---

**code.tsx**

```tsx
import React, { useEffect, useRef, useState } from 'react'
import { createRoot } from 'react-dom/client'

import {
  Module, compile, isWebGPUAvailable, Linear,
  mul, sub, mean, reshape, relu, sigmoid, concat,
  sin, cos, square,
  type Tensor, type CompiledTraining, type CompiledForward,
} from 'tensorgrad'

const IMG_W = 64
const IMG_H = 64
const N_PIXELS = IMG_W * IMG_H
const BATCH_SIZE = 1024
const L_FREQS = 8
const HIDDEN = 64

const FREQS = new Float32Array(L_FREQS)
for (let k = 0; k < L_FREQS; k++) FREQS[k] = Math.PI * Math.pow(2, k)

const GRID_COORDS = new Float32Array(N_PIXELS * 2)
for (let row = 0; row < IMG_H; row++) {
  for (let col = 0; col < IMG_W; col++) {
    const i = row * IMG_W + col
    GRID_COORDS[i * 2]     = 2 * (col + 0.5) / IMG_W - 1
    GRID_COORDS[i * 2 + 1] = 2 * (row + 0.5) / IMG_H - 1
  }
}

class NeRFTiny extends Module {
  l1 = new Linear(4 * L_FREQS, HIDDEN)
  l2 = new Linear(HIDDEN, HIDDEN)
  l3 = new Linear(HIDDEN, HIDDEN)
  l4 = new Linear(HIDDEN, 3)
}

function posEnc(coords: Tensor, freqs: Tensor): Tensor {
  const B = coords.shape[0]!
  const scaled = mul(reshape(coords, [B, 2, 1]), reshape(freqs, [1, 1, L_FREQS]))
  const sinF = reshape(sin(scaled), [B, 2 * L_FREQS])
  const cosF = reshape(cos(scaled), [B, 2 * L_FREQS])
  return concat([sinF, cosF], 1)
}

function modelFwd(m: NeRFTiny, coords: Tensor, freqs: Tensor): Tensor {
  let h = posEnc(coords, freqs)
  h = relu(m.l1.fwd(h))
  h = relu(m.l2.fwd(h))
  h = relu(m.l3.fwd(h))
  return sigmoid(m.l4.fwd(h))
}

function lossFn(
  m: NeRFTiny,
  { coords, rgb, freqs }: { coords: Tensor; rgb: Tensor; freqs: Tensor },
): Tensor {
  return mean(square(sub(modelFwd(m, coords, freqs), rgb)))
}

function predictFn(m: NeRFTiny, { coords, freqs }: { coords: Tensor; freqs: Tensor }): Tensor {
  return modelFwd(m, coords, freqs)
}

const PIXEL_SCALE = 4

function App() {
  const [status, setStatus]   = useState('initializing…')
  const [ready, setReady]     = useState(false)
  const [running, setRunning] = useState(false)

  // Long-lived (non-render) state lives in refs so updates don't trigger re-renders.
  const trainRef    = useRef<CompiledTraining<NeRFTiny> | null>(null)
  const inferRef       = useRef<CompiledForward<NeRFTiny> | null>(null)
  const targetRgbRef   = useRef<Float32Array | null>(null)
  const runningRef     = useRef(false)
  const stepRef        = useRef(0)

  // Canvas refs.
  const targetCanvasRef = useRef<HTMLCanvasElement>(null)
  const reconCanvasRef  = useRef<HTMLCanvasElement>(null)

  // Offscreen 64×64 buffer for pixel writes; drawImage upscales to visible canvas.
  const smallRef    = useRef<HTMLCanvasElement | null>(null)
  const smallCtxRef = useRef<CanvasRenderingContext2D | null>(null)
  const imgDataRef  = useRef<ImageData | null>(null)

  function rgbToCanvas(rgb: Float32Array, dstCanvas: HTMLCanvasElement) {
    const imgData = imgDataRef.current!
    const smallCtx = smallCtxRef.current!
    const small = smallRef.current!
    for (let i = 0; i < N_PIXELS; i++) {
      const off = i * 4
      imgData.data[off]     = Math.max(0, Math.min(255, Math.round(rgb[i * 3]!     * 255)))
      imgData.data[off + 1] = Math.max(0, Math.min(255, Math.round(rgb[i * 3 + 1]! * 255)))
      imgData.data[off + 2] = Math.max(0, Math.min(255, Math.round(rgb[i * 3 + 2]! * 255)))
      imgData.data[off + 3] = 255
    }
    smallCtx.putImageData(imgData, 0, 0)
    const dstCtx = dstCanvas.getContext('2d')!
    dstCtx.drawImage(small, 0, 0, dstCanvas.width, dstCanvas.height)
  }

  function nextBatch() {
    const targetRgb = targetRgbRef.current!
    const coords = new Float32Array(BATCH_SIZE * 2)
    const rgb = new Float32Array(BATCH_SIZE * 3)
    for (let b = 0; b < BATCH_SIZE; b++) {
      const idx = Math.floor(Math.random() * N_PIXELS)
      coords[b * 2]     = GRID_COORDS[idx * 2]!
      coords[b * 2 + 1] = GRID_COORDS[idx * 2 + 1]!
      rgb[b * 3]     = targetRgb[idx * 3]!
      rgb[b * 3 + 1] = targetRgb[idx * 3 + 1]!
      rgb[b * 3 + 2] = targetRgb[idx * 3 + 2]!
    }
    return { coords, rgb, freqs: FREQS }
  }

  async function renderReconstruction() {
    const infer = inferRef.current
    if (!infer) return
    const r = await infer.run({ coords: GRID_COORDS, freqs: FREQS })
    if (r.kind === 'completed') rgbToCanvas(r.output, reconCanvasRef.current!)
  }

  async function runTraining() {
    let lastRender = 0
    let lastLoss = 0
    while (runningRef.current && trainRef.current) {
      const sr = await trainRef.current.step(nextBatch())
      if (sr.kind !== 'completed') return
      lastLoss = sr.loss
      stepRef.current += 1
      if (!Number.isFinite(lastLoss)) {
        setStatus(`step ${stepRef.current}: loss is ${lastLoss} — NaN, aborting.`)
        runningRef.current = false
        setRunning(false)
        return
      }
      const now = Date.now()
      if (now - lastRender > 250) {
        lastRender = now
        setStatus(`step ${stepRef.current}  loss ${lastLoss.toFixed(5)}`)
        await renderReconstruction()
      }
      if (stepRef.current % 4 === 0) await new Promise(r => setTimeout(r, 0))
    }
  }

  async function buildGraphs() {
    setStatus('compiling…')
    const t0 = performance.now()
    const model = new NeRFTiny()
    const train = await compile({
      model,
      loss: lossFn,
      optimizer: { kind: 'adam', lr: 1e-3 },
      inputs: {
        coords: [BATCH_SIZE, 2],
        rgb:    [BATCH_SIZE, 3],
        freqs:  [L_FREQS],
      },
    })
    const infer = await train.attach({
      forward: predictFn,
      inputs: {
        coords: [N_PIXELS, 2],
        freqs:  [L_FREQS],
      },
    })
    trainRef.current = train
    inferRef.current = infer
    stepRef.current = 0
    setStatus(`compiled (${train.kernels.length} kernels, ${(performance.now() - t0).toFixed(0)} ms)`)
  }

  function makeDefaultImage(): Float32Array {
    const c = document.createElement('canvas')
    c.width = IMG_W
    c.height = IMG_H
    const ctx = c.getContext('2d')!
    const grad = ctx.createRadialGradient(IMG_W / 2, IMG_H / 2, 4, IMG_W / 2, IMG_H / 2, IMG_W * 0.7)
    grad.addColorStop(0, '#f97316')
    grad.addColorStop(1, '#1e293b')
    ctx.fillStyle = grad
    ctx.fillRect(0, 0, IMG_W, IMG_H)
    ctx.fillStyle = '#ffffff'
    ctx.font = 'bold 10px ui-sans-serif, system-ui, sans-serif'
    ctx.textAlign = 'center'
    ctx.textBaseline = 'middle'
    ctx.fillText('tensorgrad', IMG_W / 2, IMG_H / 2)
    const px = ctx.getImageData(0, 0, IMG_W, IMG_H).data
    const rgb = new Float32Array(N_PIXELS * 3)
    for (let i = 0; i < N_PIXELS; i++) {
      rgb[i * 3]     = px[i * 4]!     / 255
      rgb[i * 3 + 1] = px[i * 4 + 1]! / 255
      rgb[i * 3 + 2] = px[i * 4 + 2]! / 255
    }
    return rgb
  }

  async function loadUploadedImage(file: File): Promise<Float32Array> {
    const bitmap = await createImageBitmap(file)
    const smallCtx = smallCtxRef.current!
    smallCtx.clearRect(0, 0, IMG_W, IMG_H)
    smallCtx.drawImage(bitmap, 0, 0, IMG_W, IMG_H)
    const px = smallCtx.getImageData(0, 0, IMG_W, IMG_H).data
    const rgb = new Float32Array(N_PIXELS * 3)
    for (let i = 0; i < N_PIXELS; i++) {
      rgb[i * 3]     = px[i * 4]!     / 255
      rgb[i * 3 + 1] = px[i * 4 + 1]! / 255
      rgb[i * 3 + 2] = px[i * 4 + 2]! / 255
    }
    return rgb
  }

  async function boot() {
    if (!isWebGPUAvailable()) {
      setStatus('WebGPU not available. Try Chrome 113+ or Safari 17.4+.')
      return
    }
    const rgb = makeDefaultImage()
    targetRgbRef.current = rgb
    rgbToCanvas(rgb, targetCanvasRef.current!)
    await buildGraphs()
    await renderReconstruction()
    setReady(true)
  }

  // Mount: size canvases, build offscreen buffer, boot.
  useEffect(() => {
    const target = targetCanvasRef.current!
    const recon  = reconCanvasRef.current!
    target.width  = recon.width  = IMG_W * PIXEL_SCALE
    target.height = recon.height = IMG_H * PIXEL_SCALE
    target.getContext('2d')!.imageSmoothingEnabled = false
    recon.getContext('2d')!.imageSmoothingEnabled  = false

    const small = document.createElement('canvas')
    small.width = IMG_W
    small.height = IMG_H
    const smallCtx = small.getContext('2d')!
    smallRef.current    = small
    smallCtxRef.current = smallCtx
    imgDataRef.current  = smallCtx.createImageData(IMG_W, IMG_H)

    boot().catch((e: unknown) => {
      const msg = (e as { message?: string })?.message ?? String(e)
      setStatus(`error: ${msg}`)
      console.error(e)
    })
  }, [])

  // Event handlers.
  const onTrain = () => {
    if (!trainRef.current || runningRef.current) return
    runningRef.current = true
    setRunning(true)
    void runTraining().then(() => { setRunning(false) })
  }

  const onStop = () => {
    runningRef.current = false
    setRunning(false)
  }

  const onReset = async () => {
    const train = trainRef.current
    if (!train) return
    const wasRunning = runningRef.current
    runningRef.current = false
    setRunning(false)
    await new Promise<void>(r => setTimeout(r, 0))
    await train.reset()
    stepRef.current = 0
    await renderReconstruction()
    setStatus(`weights re-initialized (seed ${train.seed})`)
    if (wasRunning) {
      runningRef.current = true
      setRunning(true)
      void runTraining().then(() => { setRunning(false) })
    }
  }

  const onUpload = async (e: React.ChangeEvent<HTMLInputElement>) => {
    const file = e.target.files?.[0]
    if (!file) return
    const wasRunning = runningRef.current
    runningRef.current = false
    setRunning(false)
    await new Promise<void>(r => setTimeout(r, 0))
    const rgb = await loadUploadedImage(file)
    targetRgbRef.current = rgb
    rgbToCanvas(rgb, targetCanvasRef.current!)
    if (trainRef.current) {
      await trainRef.current.reset()
      stepRef.current = 0
      await renderReconstruction()
    }
    e.target.value = ''
    setStatus(`loaded ${file.name} (${rgb.length / 3} px) — weights reset`)
    if (wasRunning) {
      runningRef.current = true
      setRunning(true)
      void runTraining().then(() => { setRunning(false) })
    }
  }

  return (
    <div className="page">
      <h1>NeRF-tiny</h1>
      <p className="subtitle">
        A 4-layer MLP learns the mapping <code>(x, y) → (r, g, b)</code> for one
        64×64 image. Watch the reconstruction sharpen as the network fits the
        target. The headline trick is sinusoidal positional encoding —
        <code>sin</code> / <code>cos</code> at <em>L</em> frequency bands — without
        which a plain MLP can only fit the low-frequency part of the signal.
      </p>

      <div id="status">{status}</div>

      <div className="controls">
        <button className="primary" disabled={!ready || running} onClick={onTrain}>Train</button>
        <button disabled={!running} onClick={onStop}>Stop</button>
        <button disabled={!ready} onClick={onReset}>Reset weights</button>
        <label className="upload">
          <span>Upload image…</span>
          <input type="file" accept="image/*" disabled={!ready} onChange={onUpload} />
        </label>
      </div>

      <div className="canvases">
        <figure>
          <canvas ref={targetCanvasRef} />
          <figcaption>target</figcaption>
        </figure>
        <figure>
          <canvas ref={reconCanvasRef} />
          <figcaption>reconstruction</figcaption>
        </figure>
      </div>

      <p className="help">
        Architecture: positional encoding (<code>L = 8</code> bands, 32 features)
        → 3 × Linear(64) + ReLU → Linear(3) + sigmoid. Adam, MSE loss in pixel
        space, batches of 1024 random pixels. ~10K parameters; converges to a
        recognisable image in a few seconds.
      </p>
    </div>
  )
}

createRoot(document.getElementById('root')!).render(<App />)

```
**styles.css**

```css
:root {
  color-scheme: light;
  --bg: #ffffff;
  --fg: #1a1a1a;
  --muted: #666666;
  --border: #dddddd;
  --accent: #0d9488;
  --accent-fg: #ffffff;
  --btn-bg: #ffffff;
  --btn-border: #bbbbbb;
  --canvas-bg: #000000;
}

html[data-theme="dark"] {
  color-scheme: dark;
  --bg: #1c1c1c;
  --fg: #ececec;
  --muted: #a0a0a0;
  --border: #333333;
  --accent: #14b8a6;
  --accent-fg: #0f1f1d;
  --btn-bg: #2a2a2a;
  --btn-border: #444444;
  --canvas-bg: #000000;
}

body {
  margin: 0;
  background: var(--bg);
  color: var(--fg);
  font-family: ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
}

.page {
  max-width: 900px;
  margin: 2rem auto;
  padding: 0 1rem;
}

h1 {
  font-size: 1.4rem;
  margin-bottom: 0.25rem;
  color: var(--accent);
}

.subtitle {
  color: var(--muted);
  margin-top: 0;
  margin-bottom: 1.5rem;
}

#status {
  font-family: ui-monospace, monospace;
  font-size: 0.9rem;
  color: var(--muted);
  background: var(--btn-bg);
  border: 1px solid var(--border);
  border-radius: 6px;
  padding: 0.6rem 0.9rem;
  min-height: 1.4em;
  margin-bottom: 1rem;
}

.controls {
  display: flex;
  gap: 0.5rem;
  align-items: center;
  flex-wrap: wrap;
  margin-bottom: 1rem;
}

button, label.upload {
  font: inherit;
  padding: 0.5rem 1rem;
  font-size: 0.95rem;
  cursor: pointer;
  background: var(--btn-bg);
  color: var(--fg);
  border: 1px solid var(--btn-border);
  border-radius: 4px;
}

button.primary {
  background: var(--accent);
  color: var(--accent-fg);
  border-color: var(--accent);
}

button:disabled, label.upload:has(input:disabled) {
  opacity: 0.5;
  cursor: default;
}

label.upload input { display: none; }

.canvases {
  display: flex;
  gap: 1.5rem;
  align-items: flex-start;
  flex-wrap: wrap;
}

.canvases figure { margin: 0; }

.canvases figcaption {
  font-size: 0.85rem;
  color: var(--muted);
  margin-top: 0.4rem;
  text-align: center;
}

canvas {
  display: block;
  border: 1px solid var(--border);
  border-radius: 4px;
  background: var(--canvas-bg);
  image-rendering: pixelated;
  max-width: 100%;
  height: auto;
}

.help {
  color: var(--muted);
  font-size: 0.9rem;
  margin-top: 1rem;
}

code {
  font-family: ui-monospace, monospace;
  font-size: 0.9em;
  background: var(--border);
  padding: 0.05rem 0.3rem;
  border-radius: 3px;
}

```
**index.html**

```html
<div id="root"></div>

```
**config.json**

```json
{
  "description": "NeRF is a coordinate-network MLP learns the (x, y) → RGB mapping for one 64×64 image, using sinusoidal positional encoding so it can fit high-frequency detail.",
  "dependencies": {
    "tensorgrad": "^0.4.2",
    "react": "^19.2.6",
    "react-dom": "^19.2.6"
  }
}
```