v0.0.4
fn (%x: Tensor[(1, 4), int32]) -> Tensor[(1, 4), int32] {
  let %x1: Tensor[(1, 4), int32] = nn.batch_flatten(%x);
  %x1
}