With Einlang, you can write codes as

let x = [

    [[1.0, 2.0], [3.0, 4.0]],

    [[5.0, 6.0], [7.0, 8.0]]
];

let row_sums[..batch, i] = sum[j](x[..batch, i, j]);

let loss = sum[..batch, i](row_sums[..batch, i] * row_sums[..batch, i]);

let dloss_dx = @loss / @x;

Einlang also supports recurrence. You can write codes like

let alpha = 0.25;

let x[0] = 8.0;

let x[k in 1..6] = {

    let prev = x[k - 1];

    let loss = prev * prev;

    let g = @loss / @prev;

    prev - alpha * g
};