fill_diagonal not yet implemented

Just a feature request to implement jnp.fill_diagonal, mainly for consistency with NumPy. Though is the reason this hasn’t been implemented yet that it isn’t pure in NumPy? I don’t see an issue with having a non-pure version that returns a new array.

[jakevdp]

I see; here’s a quick implementation similar to np.fill_diagonal that should work for any batch size without reshapes or vmaps:

import jax.numpy as jnp

def fill_diagonal(a, val):
  assert a.ndim >= 2
  i, j = jnp.diag_indices(min(a.shape[-2:]))
  return a.at[..., i, j].set(val)

a = jnp.zeros((2, 3, 4, 4))

# works for scalars
a1 = fill_diagonal(a, 2)

# or for batched vectors
a2 = fill_diagonal(a, jnp.arange(24).reshape(2, 3, 4))

[hawkinsp]

My main hesitation here is that np.fill_diagonal is specified to update its input in-place. We can’t do that in JAX. We could provide an alternate API that returns an updated array, although that wouldn’t satisfy the “consistent with NumPy” constraint. What do you think?