///|
impl BackwardOp for GatherRowsBackward with fn backward(
  self,
  grad,
  _out_shape,
  inputs,
) {
  let weight_shape = inputs[0].shape
  let width = weight_shape[1]
  let weight_grad = Array::make(shape_size(weight_shape), 0.0)
  for i in 0..