Should jax.jit optimize away diagonal matrix multiplications?

[ryan112358]

From a readability standpoint it is sometimes preferable to write A @ jnp.diag(v), where A is size n x n and v is size n, even though A * v is equivalent due to broadcasting. With numpy it's clear the latter is more efficient, but one might hope that the former implementation would get compiled to the latter one with jax+jit. That doesn't appear to be the case:

@jax.jit
def mul(A, v):
  return A @ jnp.diag(v)

@jax.jit
def mul2(A, v):
  return A * v

n = 2**13
A = jnp.array(np.random.rand(n,n))
v = jnp.array(np.random.rand(n))
np.linalg.norm(mul(A, v) - mul2(A, v))

0.0

# Now that each function has been called once, we can do timing experiments
A = jnp.array(np.random.rand(n,n))
v = jnp.array(np.random.rand(n))
%timeit mul(A, v)
%timeit mul2(A, v)

1 loops, best of 5: 2.44 s per loop
10 loops, best of 5: 21.5 ms per loop

[jakevdp]

No, I don't believe the compiler has any optimization that will fuse diag plus matmul. Your best bet is probably to write the efficient version of this operation manually, as you did in mul2.

Adding such an optimization would be tricky because there is no XLA primitive for constructing a 2D diagonal matrix from a 1D vector. Consider the jaxpr for jnp.diag:

>>>  jax.make_jaxpr(jnp.diag)(jnp.arange(4))
{ lambda ; a:i32[4]. let
    b:i32[4,4] = pjit[
      jaxpr={ lambda ; c:i32[4]. let
          d:i32[4] = pad[padding_config=((0, 0, 0),)] c 0
          e:i32[4,4] = iota[dimension=0 dtype=int32 shape=(4, 4)] 
          f:i32[4,4] = add e 0
          g:i32[4,4] = iota[dimension=1 dtype=int32 shape=(4, 4)] 
          h:bool[4,4] = eq f g
          i:i32[4] = broadcast_in_dim[broadcast_dimensions=() shape=(4,)] 0
          j:i32[4,4] = pjit[
            jaxpr={ lambda ; k:bool[4,4] l:i32[4] m:i32[4]. let
                n:i32[4,4] = broadcast_in_dim[
                  broadcast_dimensions=(1,)
                  shape=(4, 4)
                ] l
                o:i32[4,4] = broadcast_in_dim[
                  broadcast_dimensions=(1,)
                  shape=(4, 4)
                ] m
                p:i32[4,4] = select_n k o n
              in (p,) }
            name=_where
          ] h d i
        in (j,) }
      name=_diag
    ] a
  in (b,) }

The compiler would have to pattern-match this sequence of ops in order to recognize it as "create diagonal matrix" and optimize it away, and that's probably not an optimization that XLA will support.

[krzysztofrusek]

@jakevdp Any chance to get spars compatible diag? The issue is that there is no sparsify rule for pad primitive.

[jakevdp]

I can't think of any way that we could make sparisify(jnp.diag) automatically do what you have in mind. You could certainly construct a diagonal sparse matrix manually if you wish, or we could think about creating a sparse.diag function that would do it for you.

[krzysztofrusek]

Having sparse.diag would be super useful to avoid boilerplate code for manual construction and I believe this is quite common pattern (aka spdiags in matlab).

I was thinking about sparsify rule for pad as this was the first error from sparsity(jnp.diag) and its easy when padding_config has only low and high makes trouble interior to provide a general rule.
Here is a gist of my tinkering but probably there is more ops not supported by sparsity in it.

[jakevdp]

Yeah, you can pick away at some of the primitives in jnp.diag, but the deeper issue is that the diagonal matrix is created from the outer equality of two dense arrays, and there's no way for sparsify to "know" in isolation that this dense-input dense-output operation is something that should be changed to dense-input sparse-output.

[krzysztofrusek]

Got it, thanks for clarification.

Btw I am working with huge Markov chains.
The transition matrix $Q$ is sparse and unnormalized, thus I want to subtract a sparse diagonal matrix whose diagonal is sum of all columns in $Q$. This is clearly sparse in sparse out operation, do you have any suggestion how to implement this efficiently under sparsify transform?