ValueError: Dimensions must be equal, but are 2 and 4 for '{{node BroadcastTo_35}} = BroadcastTo[T=DT_BOOL, Tidx=DT_INT32](All_3, Const_133)' with input shapes: [2,1], [2] and with input tensors computed as partial shapes: input[1] = [4,1].

[DarrenZhang01]

As titled, seems to be similar to issue #15 , but have not figured out the exact reason for this one.

stax.py", line 3405, in <listcomp>
    masks = [m if m is None else np.broadcast_to(
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/numpy_ops/np_array_ops.py", line 1036, in broadcast_to
    return full(shape, array)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/numpy_ops/np_array_ops.py", line 148, in full
    array_ops.broadcast_to(fill_value.data, shape))
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/gen_array_ops.py", line 854, in broadcast_to
    _, _, _op, _outputs = _op_def_library._apply_op_helper(
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/framework/op_def_library.py", line 773, in _apply_op_helper
    op = g._create_op_internal(op_type_name, inputs, dtypes=None,
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 591, in _create_op_internal
    return super(FuncGraph, self)._create_op_internal(  # pylint: disable=protected-access
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 3510, in _create_op_internal
    ret = Operation(
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 2001, in __init__
    self._c_op = _create_c_op(self._graph, node_def, inputs,
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 1842, in _create_c_op
    raise ValueError(str(e))
ValueError: Dimensions must be equal, but are 2 and 4 for '{{node BroadcastTo_35}} = BroadcastTo[T=DT_BOOL, Tidx=DT_INT32](All_3, Const_133)' with input shapes: [2,1], [2] and with input tensors computed as partial shapes: input[1] = [4,1].

[DarrenZhang01]

Solved in 966e96b

As a note to issue google/neural-tangents#56, after experimenting, it turns out that the basic idea is: Assume we have nested structures parallel{serial1, serial2, serial3}, and the goal is to construct the Neural Tangent Kernel, the kernel functions for the inside layers will be first constructed, i.e., serial1, 2 and 3. And then the outside layer - parallel will construct its kernel on top of the kernels of serial1, 2 and 3. The point is that when constructing the kernel, the init_fun of each layer will be called in setting up shape attributes of the Kernel object. And that is why the init_fun of serial1, 2 and 3 are called before the parallel block.