Add a JAX implementation for `FillDiagonal`

aesara/link/jax/dispatch.py

@@ -956,20 +956,11 @@ def jax_funcify_Bartlett(op, **kwargs):
 
 @jax_funcify.register(FillDiagonal)
 def jax_funcify_FillDiagonal(op, **kwargs):
+    def filldiagonal(value, diagonal):
+        i, j = jnp.diag_indices(min(value.shape[-2:]))
+        return value.at[..., i, j].set(diagonal)
 
-    # def filldiagonal(a, val):
-    #     if a.ndim == 2:
-    #         step = a.shape[1] + 1
-    #         end = a.shape[1] * a.shape[1]
-    #         a.flat[:end:step] = val
-    #     else:
-    #         jnp.fill_diagonal(a, val)
-    #
-    #     return a
-    #
-    # return filldiagonal
-
-    raise NotImplementedError("flatiter not implemented in JAX")
+    return filldiagonal
 
 
 @jax_funcify.register(FillDiagonalOffset)

tests/link/test_jax.py

@@ -1205,7 +1205,6 @@ def test_extra_ops():
 
     c = at.as_tensor(5)
 
-    with pytest.raises(NotImplementedError):
     out = at_extra_ops.fill_diagonal(a, c)
     fgraph = FunctionGraph([a], [out])
     compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])