<p dir="auto">Currently, the following script</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">scipy</span>.<span class="pl-s1">linalg</span> <span class="pl-k">import</span> <span class="pl-s1">solve_triangular</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>

<span class="pl-s1">dim</span> <span class="pl-c1">=</span> <span class="pl-c1">130</span>
<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">broadcast_to</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">eye</span>(<span class="pl-s1">dim</span>), (<span class="pl-c1">2</span>, <span class="pl-s1">dim</span>, <span class="pl-s1">dim</span>))
<span class="pl-en">solve_triangular</span>(<span class="pl-s1">x</span>, <span class="pl-s1">x</span>)  <span class="pl-c"># error</span></pre></div>
<p dir="auto">raises the issue: <code class="notranslate">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 { }.: </code></p>
<p dir="auto">This does not fail when</p>
<ul dir="auto">
<li><code class="notranslate">dim &lt;= 128</code></li>
<li>no batching (<code class="notranslate">(2,)</code> in the above script)</li>
</ul>
<p dir="auto">I am trying to apply <code class="notranslate">vmap</code> to a matrix solve, but it errors out on matrices larger than 128 x 128.</p>
<p dir="auto">Here's an example:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>

<span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">vmap</span>

<span class="pl-k">def</span> <span class="pl-en">vmap_solve</span>(<span class="pl-v">N</span>):
  <span class="pl-v">A</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">expand_dims</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">eye</span>(<span class="pl-v">N</span>), <span class="pl-c1">0</span>)
  <span class="pl-v">B</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-v">N</span>)
  <span class="pl-k">return</span> <span class="pl-en">vmap</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span>.<span class="pl-s1">linalg</span>.<span class="pl-s1">solve</span>, <span class="pl-s1">in_axes</span><span class="pl-c1">=</span>(<span class="pl-c1">0</span>, <span class="pl-c1">None</span>))(<span class="pl-v">A</span>, <span class="pl-v">B</span>)</pre></div>
<p dir="auto"><code class="notranslate">vmap_solve(128)</code> gives a vector of ones, as expected.</p>
<p dir="auto"><code class="notranslate">vmap_solve(129)</code> produces the following error:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="RuntimeError                              Traceback (most recent call last)

&lt;ipython-input-16-31a363cce168&gt; in &lt;module&gt;()
----&gt; 1 print(vmap_solve(129))

16 frames

&lt;ipython-input-9-fc01e1c69adf&gt; in vmap_solve(N)
      2   A = jnp.expand_dims(jnp.eye(N), 0)
      3   B = jnp.ones(N)
----&gt; 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(&quot;vmap out_axes&quot;, out_tree(),
   1232                                                    out_axes),
-&gt; 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)
---&gt; 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:
--&gt; 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__,
--&gt; 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):
-&gt; 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):
-&gt; 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):
-&gt; 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)
--&gt; 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):
-&gt; 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):
-&gt; 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):
-&gt; 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):
--&gt; 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,
--&gt; 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:
--&gt; 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
--&gt; 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
--&gt; 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 { }.: "><pre class="notranslate"><code class="notranslate">RuntimeError                              Traceback (most recent call last)

&lt;ipython-input-16-31a363cce168&gt; in &lt;module&gt;()
----&gt; 1 print(vmap_solve(129))

16 frames

&lt;ipython-input-9-fc01e1c69adf&gt; in vmap_solve(N)
      2   A = jnp.expand_dims(jnp.eye(N), 0)
      3   B = jnp.ones(N)
----&gt; 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),
-&gt; 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)
---&gt; 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:
--&gt; 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__,
--&gt; 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):
-&gt; 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):
-&gt; 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):
-&gt; 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)
--&gt; 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):
-&gt; 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):
-&gt; 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):
-&gt; 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):
--&gt; 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,
--&gt; 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:
--&gt; 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
--&gt; 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
--&gt; 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 { }.: 
</code></pre></div>
<p dir="auto">Any ideas on why this is happening and how it can be fixed? Thank you.</p>