Inspect how one program moves through the Zigrad pipeline.

entry: train_step
zxpr loss {
  ; params (8)
  a: 784x128<f32>
  b: 128<f32>
  c: 128x64<f32>
  d: 64<f32>
  e: 64x10<f32>
  f: 10<f32>
  g: 64x784<f32>
  h: 64x10<f32>
  ; body (4 ops)
  let
    i: 64x128<f32> = dot[contracting=([1], [0]), K=784](g, a)  ; vjp
    j: 64x128<f32> = broadcast_in_dim[[128] -> [64, 128], dims=[1]](b)  ; vjp
    k: 64x128<f32> = add(i, j)  ; vjp
    l: f32 = reduce_sum[axes=[0, 1]](k)  ; vjp
  in l
}

zxpr train_step {
  ; params (8)
  a: 784x128<f32>
  b: 128<f32>
  c: 128x64<f32>
  d: 64<f32>
  e: 64x10<f32>
  f: 10<f32>
  g: 64x784<f32>
  h: 64x10<f32>
  ; body (3 ops)
  let
    i: f32, j: 784x128<f32>, k: 128<f32> = call[callee="loss_vjp"](a, b, c, d, e, f, g, h)
    l: f32 = literal[0.01]()
    m: 784x128<f32> = broadcast_in_dim[[] -> [784, 128], dims=[]](l)  ; vjp
  in i, j, k
}

// Lowered: StableHLO -> XLA/PJRT
func.func @matmul_relu(%arg0: tensor<512x512xf32>,
                       %arg1: tensor<512x512xf32>) -> tensor<512x512xf32> {
  %0 = stablehlo.dot_general %arg0, %arg1,
       contracting_dims = [1] x [0] : (tensor<512x512xf32>,
       tensor<512x512xf32>) -> tensor<512x512xf32>
  %1 = stablehlo.maximum %0, zeroes : tensor<512x512xf32>
  return %1
}