Method slot_map

backends/mlx/builder/program_builder.py:200–222

(self, tree)

        return self.slot_map(node.kwargs)

    def slot_map(self, tree):
        leaves, spec = pytree.tree_flatten(tree)
        new_leaves = []
        for a in leaves:
            if isinstance(a, Node):
                # Use make_or_get_slots which handles both single and multi-output nodes.
                # For single-output nodes, returns a 1-tuple; for multi-output, returns n-tuple.
                # We unwrap single-element tuples for convenience.
                slots = self.make_or_get_slots(a)
                if len(slots) == 1:
                    new_leaves.append(slots[0])
                else:
                    new_leaves.append(slots)
            else:
                new_leaves.append(a)

        for a in new_leaves:
            if isinstance(a, Slot):
                assert self.slot_manager.is_alive(
                    a
                ), f"Slot {a} is not alive; it was either already freed or never created"

        return pytree.tree_unflatten(new_leaves, spec)

    def make_or_get_slots(
        self, node: Node, id_space: IdSpace = IdSpace.Temp