vmap behaving unexpected with function that inverts singular matrix

[bjeurissen]

Description

I ran into unexpected behavior when using vmap:

import jax.numpy as jnp
from jax import vmap, __version__
import platform
print(platform.platform())
print(platform.python_version())
print(__version__)

def myfun(x):
    return jnp.linalg.inv(x.T@x)

myfun_many = vmap(myfun, in_axes=2, out_axes=2)

print('test1')
x1 = jnp.array([[1.0,3.0],[-5.0,1.0]])
y1 = myfun(x1)
print(y1)

x2 = jnp.stack((x1,x1,x1),2)
assert (x2[:,:,0]==x1).all()
assert (x2[:,:,1]==x1).all()
assert (x2[:,:,2]==x1).all()
y2 = myfun_many(x2)
print(y2[:,:,0]) # this is the same as y1
print(y2[:,:,1]) # this is the same as y1
print(y2[:,:,2]) # this is the same as y1


print('test2')
x1 = jnp.array([[1.0,1.0],[1.0,1.0]])
y1 = myfun(x1)
print(y1)

x2 = jnp.stack((x1,x1,x1),2)
assert (x2[:,:,0]==x1).all()
assert (x2[:,:,1]==x1).all()
assert (x2[:,:,2]==x1).all()
y2 = myfun_many(x2)
print(y2[:,:,0]) # this is not the same as y1!
print(y2[:,:,1]) # this is not the same as y1!
print(y2[:,:,2]) # this is not the same as y1!

which outputs:

macOS-13.3-arm64-arm-64bit
3.11.1
0.4.8
test1
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
test2
[[ inf -inf]
 [-inf  inf]]
[[ 1.5 -1. ]
 [-1.   1. ]]
[[ 1.5 -1. ]
 [-1.   1. ]]
[[ 1.5 -1. ]
 [-1.   1. ]]

In the first test, one can appreciate that vmap correctly vectorized myfun when x.T@x can be inverted.

However, when x.T@x cannot be inverted, myfun "correctly" returns Infs, whereas vmap of myfun returns something else.

What jax/jaxlib version are you using?

jax 0.4.8

Which accelerator(s) are you using?

CPU

Additional system info

Python 3.11.1, macOS-13.3-arm64-arm-64bit

NVIDIA GPU info

No response

[alonfnt]

I can reproduce this on CPU. On GPU it works correctly.
CPU

Linux-5.19.0-41-generic-x86_64-with-glibc2.35
3.10.6
0.4.8
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[CpuDevice(id=0)]
test1
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
test2
[[ inf -inf]
 [-inf  inf]]
[[ 1.5 -1. ]
 [-1.   1. ]]
[[ 1.5 -1. ]
 [-1.   1. ]]
[[ 1.5 -1. ]
 [-1.   1. ]]

GPU:

Linux-5.19.0-41-generic-x86_64-with-glibc2.35
3.10.6
0.4.8
[StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0)]
test1
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
[[0.0390625  0.0078125 ]
 [0.0078125  0.10156249]]
test2
[[ inf -inf]
 [-inf  inf]]
[[ inf -inf]
 [-inf  inf]]
[[ inf -inf]
 [-inf  inf]]
[[ inf -inf]
 [-inf  inf]]

[jakevdp]

Shorter repro:

import jax
import jax.numpy as jnp

x = jnp.ones((1, 2, 2))
print(jnp.linalg.inv(x))
# [[[ inf -inf]
#   [-inf  inf]]]

print(jax.vmap(jnp.linalg.inv)(x))
# [[[ 2. -1.]
#   [-1.  1.]]]

It probably has something to do with the custom_linear_solve batching rule. I'm going to dig into it.

[jakevdp]

I think I'm getting closer – this looks like it somehow comes from the batching rule of jax.lax.triangular_solve:

import jax

def solve(x, y):
  return jax.lax.linalg.triangular_solve(x, y, left_side=True)

x = jnp.array([[1., 1.], [1., 0.]])
y = jnp.array([[1.], [0.]])

print(solve(x, y))
# [[nan]
#  [nan]]

print(jax.vmap(solve)(x[None], y[None])[0])
# [[1.]
#  [0.]]

[jakevdp]

This is still an issue as of JAX v0.4.26.