def path_matrix(self) -> pz.nx.NamedArray:
    """Builds a matrix that maps "from" indices to their "to" indices."""
    result = [[0 for _ in self.world_ordering] for _ in self.world_ordering]
    assert len(self.taking) == len(self.world_ordering)
    assert set(self.taking.keys()) == set(self.world_ordering)
    for dest, connections in self.taking.items():
      if isinstance(connections, From):
        connections = (connections,)
      for connection in connections:
        from_ix = self.world_ordering.index(connection.source)
        to_ix = self.world_ordering.index(dest)
        result[to_ix][from_ix] += connection.weight
    # Allows the weights to be named arrays by `nmap`ing the array constructor.
    # This works as long as each weight had an empty positional shape.
    return pz.nx.nmap(jnp.array)(result)
Builds a matrix that maps "from" indices to their "to" indices.
penzai/deprecated/v1/toolshed/model_rewiring.py