Multimer colab erroring out in JAX with 'incompatable shapes for broadcasting ...'

Today, a colleague tried to run a multimer model on google colab with 2 protein and an unchecked the `run_relax` option.

It errors out in cell 5 with:
```python
#5. Run Alphafold and download prediction

Running model_5_multimer_v3: 100%
6/6 [elapsed: 11:54 remaining: 00:00]
Warning: Running without the relaxation stage.
---------------------------------------------------------------------------
UnfilteredStackTrace Traceback (most recent call last)
in
123 break
--> 124 banded_b_factors = np.array(banded_b_factors)[:, None] * final_atom_mask
125 to_visualize_pdb = utils.overwrite_b_factors(relaxed_pdb, banded_b_factors)

23 frames
/usr/local/lib/python3.8/dist-packages/jax/_src/numpy/lax_numpy.py in deferring_binary_op(self, other)
5070 if isinstance(other, _accepted_binop_types):
-> 5071 return binary_op(*args)
5072 if isinstance(other, _rejected_binop_types):

/usr/local/lib/python3.8/dist-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
162 try:
--> 163 return fun(*args, **kwargs)
164 except Exception as e:

/usr/local/lib/python3.8/dist-packages/jax/_src/pjit.py in cache_miss(*args, **kwargs)
236 def cache_miss(*args, **kwargs):
--> 237 outs, out_flat, out_tree, args_flat = _python_pjit_helper(
238 fun, infer_params_fn, *args, **kwargs)

/usr/local/lib/python3.8/dist-packages/jax/_src/pjit.py in _python_pjit_helper(fun, infer_params_fn, *args, **kwargs)
179 def _python_pjit_helper(fun, infer_params_fn, *args, **kwargs):
--> 180 args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
181 *args, **kwargs)

/usr/local/lib/python3.8/dist-packages/jax/_src/api.py in infer_params(*args, **kwargs)
442 inline=inline, resource_env=None)
--> 443 return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
444

/usr/local/lib/python3.8/dist-packages/jax/_src/pjit.py in common_infer_params(***failed resolving arguments***)
519
--> 520 jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
521 flat_fun, hashable_pytree(out_shardings), global_in_avals,

/usr/local/lib/python3.8/dist-packages/jax/_src/linear_util.py in memoized_fun(fun, *args)
300 else:
--> 301 ans = call(fun, *args)
302 cache[key] = (ans, fun.stores)

/usr/local/lib/python3.8/dist-packages/jax/_src/pjit.py in _pjit_jaxpr(fun, out_shardings_thunk, global_in_avals, out_tree, api_name)
931 event=dispatch.JAXPR_TRACE_EVENT):
--> 932 jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
933 fun, global_in_avals, debug_info=pe.debug_info_final(fun, api_name))

/usr/local/lib/python3.8/dist-packages/jax/_src/profiler.py in wrapper(*args, **kwargs)
313 with TraceAnnotation(name, **decorator_kwargs):
--> 314 return func(*args, **kwargs)
315 return wrapper

/usr/local/lib/python3.8/dist-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr_dynamic(fun, in_avals, debug_info, keep_inputs)
1984 main.jaxpr_stack = () # type: ignore
-> 1985 jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
1986 fun, main, in_avals, keep_inputs=keep_inputs, debug_info=debug_info)

/usr/local/lib/python3.8/dist-packages/jax/interpreters/partial_eval.py in trace_to_subjaxpr_dynamic(fun, main, in_avals, keep_inputs, debug_info)
2001 in_tracers_ = [t for t, keep in zip(in_tracers, keep_inputs) if keep]
-> 2002 ans = fun.call_wrapped(*in_tracers_)
2003 out_tracers = map(trace.full_raise, ans)

/usr/local/lib/python3.8/dist-packages/jax/_src/linear_util.py in call_wrapped(self, *args, **kwargs)
164 try:
--> 165 ans = self.f(*args, **dict(self.params, **kwargs))
166 except:

/usr/local/lib/python3.8/dist-packages/jax/_src/numpy/ufuncs.py in fn(x1, x2)
97 x1, x2 = _promote_args(numpy_fn.__name__, x1, x2)
---> 98 return lax_fn(x1, x2) if x1.dtype != np.bool_ else bool_lax_fn(x1, x2)
99 fn.__qualname__ = f"jax.numpy.{numpy_fn.__name__}"

/usr/local/lib/python3.8/dist-packages/jax/_src/lax/lax.py in mul(x, y)
462 r"""Elementwise multiplication: :math:`x \times y`."""
--> 463 return mul_p.bind(x, y)
464

/usr/local/lib/python3.8/dist-packages/jax/_src/core.py in bind(self, *args, **params)
342 all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args
--> 343 return self.bind_with_trace(find_top_trace(args), args, params)
344

/usr/local/lib/python3.8/dist-packages/jax/_src/core.py in bind_with_trace(self, trace, args, params)
345 def bind_with_trace(self, trace, args, params):
--> 346 out = trace.process_primitive(self, map(trace.full_raise, args), params)
347 return map(full_lower, out) if self.multiple_results else full_lower(out)

/usr/local/lib/python3.8/dist-packages/jax/interpreters/partial_eval.py in process_primitive(self, primitive, tracers, params)
1720 return custom_staging_rules[primitive](self, *tracers, **params)
-> 1721 return self.default_process_primitive(primitive, tracers, params)
1722

/usr/local/lib/python3.8/dist-packages/jax/interpreters/partial_eval.py in default_process_primitive(self, primitive, tracers, params)
1724 avals = [t.aval for t in tracers]
-> 1725 out_avals, effects = primitive.abstract_eval(*avals, **params)
1726 out_avals = [out_avals] if not primitive.multiple_results else out_avals

/usr/local/lib/python3.8/dist-packages/jax/_src/core.py in abstract_eval_(*args, **kwargs)
378 def abstract_eval_(*args, **kwargs):
--> 379 return abstract_eval(*args, **kwargs), no_effects
380 return abstract_eval_

/usr/local/lib/python3.8/dist-packages/jax/_src/lax/utils.py in standard_abstract_eval(prim, shape_rule, dtype_rule, weak_type_rule, named_shape_rule, *avals, **kwargs)
65 elif least_specialized is core.ShapedArray:
---> 66 return core.ShapedArray(shape_rule(*avals, **kwargs),
67 dtype_rule(*avals, **kwargs), weak_type=weak_type,

/usr/local/lib/python3.8/dist-packages/jax/_src/lax/lax.py in broadcasting_shape_rule(name, *avals)
1595 else:
-> 1596 raise TypeError(f'{name} got incompatible shapes for broadcasting: '
1597 f'{", ".join(map(str, map(tuple, shapes)))}.')

UnfilteredStackTrace: TypeError: mul got incompatible shapes for broadcasting: (0, 1), (652, 37).

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

TypeError Traceback (most recent call last)
in
122 banded_b_factors.append(idx)
123 break
--> 124 banded_b_factors = np.array(banded_b_factors)[:, None] * final_atom_mask
125 to_visualize_pdb = utils.overwrite_b_factors(relaxed_pdb, banded_b_factors)
126

/usr/local/lib/python3.8/dist-packages/jax/_src/numpy/lax_numpy.py in deferring_binary_op(self, other)
5069 args = (other, self) if swap else (self, other)
5070 if isinstance(other, _accepted_binop_types):
-> 5071 return binary_op(*args)
5072 if isinstance(other, _rejected_binop_types):
5073 raise TypeError(f"unsupported operand type(s) for {opchar}: "

/usr/local/lib/python3.8/dist-packages/jax/_src/numpy/ufuncs.py in fn(x1, x2)
96 def fn(x1, x2, /):
97 x1, x2 = _promote_args(numpy_fn.__name__, x1, x2)
---> 98 return lax_fn(x1, x2) if x1.dtype != np.bool_ else bool_lax_fn(x1, x2)
99 fn.__qualname__ = f"jax.numpy.{numpy_fn.__name__}"
100 fn = jit(fn, inline=True)

TypeError: mul got incompatible shapes for broadcasting: (0, 1), (652, 37).```

```

Do you have any clue how to prevent this error? (We have used the multimer models successfully in the past)