Issue 1:
Currently, the following script

from jax.scipy.linalg import solve_triangular
import jax.numpy as jnp

dim = 130
x = jnp.broadcast_to(jnp.eye(dim), (2, dim, dim))
solve_triangular(x, x) # error

raises the issue: RuntimeError: Invalid argument: The rank of the operand and the padding configuration do not match: f32[126,126] vs dimensions { } dimensions { edge_padding_low: 2 } dimensions { }.:

This does not fail when

dim <= 128

no batching ((2,) in the above script)

Issue 2:
I am trying to apply vmap to a matrix solve, but it errors out on matrices larger than 128 x 128.

Here's an example:

import jax
import jax.numpy as jnp

from jax import vmap

def vmap_solve(N):
A = jnp.expand_dims(jnp.eye(N), 0)
B = jnp.ones(N)
return vmap(jax.numpy.linalg.solve, in_axes=(0, None))(A, B)

vmap_solve(128) gives a vector of ones, as expected.

vmap_solve(129) produces the following error:

RuntimeError Traceback (most recent call last)

<ipython-input-16-31a363cce168> in <module>()
----> 1 print(vmap_solve(129))

16 frames

<ipython-input-9-fc01e1c69adf> in vmap_solve(N)
2 A = jnp.expand_dims(jnp.eye(N), 0)
3 B = jnp.ones(N)
----> 4 return vmap(jax.numpy.linalg.solve, in_axes=(0, None))(A, B)

/usr/local/lib/python3.6/dist-packages/jax/api.py in batched_fun(*args)
1231 lambda: flatten_axes("vmap out_axes", out_tree(),
1232 out_axes),
-> 1233 axis_name=axis_name)
1234 return tree_unflatten(out_tree(), out_flat)
1235

/usr/local/lib/python3.6/dist-packages/jax/interpreters/batching.py in batch(fun, in_vals, in_dims, out_dim_dests, axis_name)
34 # executes a batched version of `fun` following out_dim_dests
35 batched_fun = batch_fun(fun, in_dims, out_dim_dests, axis_name=axis_name)
---> 36 return batched_fun.call_wrapped(*in_vals)
37
38 @lu.transformation_with_aux

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
154
155 try:
--> 156 ans = self.f(*args, **dict(self.params, **kwargs))
157 except:
158 # Some transformations yield from inside context managers, so we have to

/usr/local/lib/python3.6/dist-packages/jax/api.py in f_jitted(*args, **kwargs)
217 backend=backend,
218 name=flat_fun.__name__,
--> 219 donated_invars=donated_invars)
220 return tree_unflatten(out_tree(), out)
221

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, fun, *args, **params)
1172
1173 def bind(self, fun, *args, **params):
-> 1174 return call_bind(self, fun, *args, **params)
1175
1176 def process(self, trace, fun, tracers, params):

/usr/local/lib/python3.6/dist-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
1163 tracers = map(top_trace.full_raise, args)
1164 with maybe_new_sublevel(top_trace):
-> 1165 outs = primitive.process(top_trace, fun, tracers, params)
1166 return map(full_lower, apply_todos(env_trace_todo(), outs))
1167

/usr/local/lib/python3.6/dist-packages/jax/core.py in process(self, trace, fun, tracers, params)
1175
1176 def process(self, trace, fun, tracers, params):
-> 1177 return trace.process_call(self, fun, tracers, params)
1178
1179 def post_process(self, trace, out_tracers, params):

/usr/local/lib/python3.6/dist-packages/jax/interpreters/batching.py in process_call(self, call_primitive, f, tracers, params)
169 else:
170 f, dims_out = batch_subtrace(f, self.main, dims)
--> 171 vals_out = call_primitive.bind(f, *vals, **params)
172 return [BatchTracer(self, v, d) for v, d in zip(vals_out, dims_out())]
173

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, fun, *args, **params)
1172
1173 def bind(self, fun, *args, **params):
-> 1174 return call_bind(self, fun, *args, **params)
1175
1176 def process(self, trace, fun, tracers, params):

/usr/local/lib/python3.6/dist-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
1163 tracers = map(top_trace.full_raise, args)
1164 with maybe_new_sublevel(top_trace):
-> 1165 outs = primitive.process(top_trace, fun, tracers, params)
1166 return map(full_lower, apply_todos(env_trace_todo(), outs))
1167

/usr/local/lib/python3.6/dist-packages/jax/core.py in process(self, trace, fun, tracers, params)
1175
1176 def process(self, trace, fun, tracers, params):
-> 1177 return trace.process_call(self, fun, tracers, params)
1178
1179 def post_process(self, trace, out_tracers, params):

/usr/local/lib/python3.6/dist-packages/jax/core.py in process_call(self, primitive, f, tracers, params)
574
575 def process_call(self, primitive, f, tracers, params):
--> 576 return primitive.impl(f, *tracers, **params)
577 process_map = process_call
578

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_call_impl(fun, device, backend, name, donated_invars, *args)
555 def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name, donated_invars):
556 compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
--> 557 *unsafe_map(arg_spec, args))
558 try:
559 return compiled_fun(*args)

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in memoized_fun(fun, *args)
245 fun.populate_stores(stores)
246 else:
--> 247 ans = call(fun, *args)
248 cache[key] = (ans, fun.stores)
249

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
701 device_assignment=(device.id,) if device else None)
702 options.parameter_is_tupled_arguments = tuple_args
--> 703 compiled = backend_compile(backend, built, options)
704 if nreps == 1:
705 return partial(_execute_compiled, compiled, out_avals, result_handlers)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in backend_compile(backend, built_c, options)
342 # we use a separate function call to ensure that XLA compilation appears
343 # separately in Python profiling results
--> 344 return backend.compile(built_c, compile_options=options)
345
346 def _execute_compiled_primitive(prim, compiled, result_handler, *args):

RuntimeError: Invalid argument: The rank of the operand and the padding configuration do not match: f32[127,127] vs dimensions { } dimensions { edge_padding_low: 1 } dimensions { }.:

Any ideas on why this is happening and how it can be fixed? Thank you.