add missing at functions

[davidkoski]

• fixes add missing at() function #188
• also fixes bug with array[idx] += 3 that would not assign back to array

[davidkoski]

• needs documentation

[davidkoski]

• needs documentation

[davidkoski]

This is the pattern on the python side:

mx::array mlx_subtract_item(
    const mx::array& src,
    const nb::object& obj,
    const ScalarOrArray& v) {
  auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
  if (indices.size() > 0) {
    return scatter_add(src, indices, -updates, axes);
  } else {
    return src - updates;
  }
}

This is the same but using the mlx-c API.

[davidkoski]

Found in unit tests:

        // this references each index twice
        let idx = MLXArray([0, 1, 0, 1])

        // similar to above -- we can only observe one assignment, so we just get a +1
        // note: there was a bug in the += operator where the lhs was not inout and
        // this was producing [0, 0]
        let a1 = MLXArray([0, 0])
        a1[idx] += 1
        assertEqual(a1, MLXArray([1, 1]))

Without the inout it seems that:

array[idx] += 2

is roughly equivalent to:

tmp = array[idx]
tmp += 2

[davidkoski]

Because of the inout change this needs to be var now

[angeloskath]

Looks great!

[angeloskath]

Total nitpick but I would use reciprocal here to avoid the extra array and broadcast op (although likely negligible).

[davidkoski]

• use reciprocal