`ValueError: Dimensions must be equal, but are ...`

[DarrenZhang01]

stax.py", line 1216, in mask_fn
    mask = np.broadcast_to(mask, input_shape)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/trax/tf_numpy/numpy_impl/array_ops.py", line 1237, in broadcast_to
    return full(shape, array)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/trax/tf_numpy/numpy_impl/array_ops.py", line 221, in full
    return arrays_lib.tensor_to_ndarray(tf.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 829, 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 742, 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 593, 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 3319, 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 1816, 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 1657, in _create_c_op
    raise ValueError(str(e))
ValueError: Dimensions must be equal, but are 2 and 4 for '{{node BroadcastTo_1}} = BroadcastTo[T=DT_BOOL, Tidx=DT_INT32](cond_1/Identity, BroadcastTo_1/shape)' with input shapes: [2,210], [2] and with input tensors computed as partial shapes: input[1] = [4,210].

[DarrenZhang01]

Solved in 3a7e6e3