Dimensions do not match

[DarrenZhang01]

As titled.

utils.py", line 341, in outer_prod
    return prod_op(x, y)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/trax/tf_numpy/numpy_impl/math_ops.py", line 79, in multiply
    return _bin_op(mul_or_and, x1, x2)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/trax/tf_numpy/numpy_impl/math_ops.py", line 54, in _bin_op
    return utils.tensor_to_ndarray(tf_fun(a.data, b.data))
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/trax/tf_numpy/numpy_impl/math_ops.py", line 78, in mul_or_and
    return tf.multiply(x1, x2)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py", line 180, in wrapper
    return target(*args, **kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 381, in multiply
    return gen_math_ops.mul(x, y, name)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py", line 6089, in mul
    _ops.raise_from_not_ok_status(e, name)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 6653, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.UnimplementedError: Broadcast between [5,1,8,1,8,1] and [1,5,1,8,1,8] is not supported yet. [Op:Mul]

[DarrenZhang01]

google/neural-tangents#54