Broadcast array on specific indices only

[jstpierref]

While doing a convolution with lax.conv_dimension_numbers that simply sharpens the input array, I had to manually build the kernel with loops to make sure the kernel is non-null only on channels dimensions diagonal. It is to make sure the transformations are kept seperate between channels.

kernel = jnp.zeros((3, 3, 3, 3), dtype=jnp.float32)
kernel += jnp.array([
    [-1, -1, -1],
    [-1, 16, -1],
    [-1, -1, -1],
], dtype=jnp.float32)[:, :, jnp.newaxis, jnp.newaxis]/8.0

for i in range(kernel.shape[0]):
    for j in range(kernel.shape[1]):
        kernel = kernel.at[i,j,:,:].set(jnp.eye(kernel.shape[2])*kernel[i,j,0,0])

It's nothing critical, but I'm looking for ideas to achieve this in functional ways. I tried things with einsum but examples I found in numpy were not immutable.

If that matters, I'm on Apple M1 Pro running with jax-metal Apple plugin.

[soraros]

Try these:

A = jnp.array([[-1, -1, -1], [-1, 16, -1], [-1, -1, -1]]) / 8.0
I = jnp.eye(3)
k0 = jnp.tensordot(A, I, axes=((), ()))
k1 = jnp.einsum('ij,kl->ijkl', A, I)

[jstpierref]

Solid thanks!