Skip to content

Commit 281d9b6

Browse files
ghamarianghamarian
committed
A midway fix for tensor transform app
1 parent 82c72c1 commit 281d9b6

2 files changed

Lines changed: 132 additions & 42 deletions

File tree

‎tensor_transform_app.py‎

Lines changed: 108 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -506,6 +506,11 @@ def build_transformation_graph_from_pytensor(descriptors, variables):
506506

507507
# Process each descriptor as a separate stage
508508
actual_stage_num = 1 # Track actual stage numbers starting from 1
509+
last_stage_outputs = prev_stage_output_nodes.copy() # Track only the most recent stage's outputs
510+
511+
# For multi-stage descriptors, track the dimension mapping from previous stage
512+
previous_stage_dim_map = {} # Maps logical dimension index to node ID
513+
509514
for stage_idx, tensor_desc in enumerate(pytensor_descriptors):
510515
print(f"DEBUG: Processing stage_idx={stage_idx}")
511516
stage_output_nodes = {}
@@ -606,9 +611,10 @@ def build_transformation_graph_from_pytensor(descriptors, variables):
606611
lower_indices = new_lower_idss[i]
607612
upper_indices = new_upper_idss[i]
608613

609-
print(f"DEBUG: Transform {i} - upper_indices (inputs): {upper_indices}, lower_indices (outputs): {lower_indices}")
610-
print(f"DEBUG: Transform {i} type: {transform.__class__.__name__}")
611-
print(f"DEBUG: prev_stage_output_nodes keys: {list(prev_stage_output_nodes.keys())}")
614+
print(f"DEBUG: Stage {stage_idx}, Transform {i} type: {transform.__class__.__name__}")
615+
print(f"DEBUG: lower_indices (dimension inputs): {lower_indices}")
616+
print(f"DEBUG: upper_indices (dimension outputs): {upper_indices}")
617+
print(f"DEBUG: Available from prev stage: {list(last_stage_outputs.keys()) if stage_idx > 0 else 'Initial inputs'}")
612618

613619
# Extra debug for UnmergeTransform in naive packed
614620
if (isinstance(transform, pytensor.tensor_descriptor.UnmergeTransform) and
@@ -737,21 +743,31 @@ def build_transformation_graph_from_pytensor(descriptors, variables):
737743
# Special handling for UnmergeTransform in naive packed descriptors
738744
# The output indices use hidden dimension IDs that may not start from 0
739745
if (isinstance(transform, pytensor.tensor_descriptor.UnmergeTransform) and
740-
stage_idx == 0 and is_first_naive_packed and
741-
len(descriptors) == 1): # Only for simple naive packed without subsequent transforms
742-
print(f"DEBUG: UnmergeTransform in simple naive packed - remapping indices")
746+
stage_idx == 0 and is_first_naive_packed):
747+
print(f"DEBUG: UnmergeTransform in naive packed - remapping indices")
743748
print(f"DEBUG: Original output_indices: {output_indices}")
744-
# For simple naive packed, map outputs to logical indices 0,1,2,3...
745-
if len(output_formulas) == num_logical_dims:
746-
filtered_output_indices = list(range(num_logical_dims))
749+
# For naive packed (even in multi-stage), map outputs to logical indices 0,1,2,3...
750+
# This ensures the next stage can find these outputs
751+
if hasattr(transform, 'lengths') and transform.lengths:
752+
num_outputs = len(transform.lengths)
753+
filtered_output_indices = list(range(num_outputs))
747754
print(f"DEBUG: Remapped to logical indices: {filtered_output_indices}")
748755

749756
# Create output nodes for each formula within the stage cluster
750757
for j, output_idx in enumerate(filtered_output_indices):
751758
if j < len(output_formulas):
752759
# Use the actual stage number for node naming
753760
node_id = f"s{actual_stage_num}_t{i}_d{output_idx}"
754-
stage_output_nodes[output_idx] = node_id
761+
762+
# For multi-stage examples with UnmergeTransform in stage 0,
763+
# use sequential logical indices to ensure proper mapping to next stage
764+
if (stage_idx == 0 and len(pytensor_descriptors) > 1 and
765+
isinstance(transform, pytensor.tensor_descriptor.UnmergeTransform)):
766+
# Use sequential index j instead of output_idx for storage key
767+
stage_output_nodes[j] = node_id
768+
print(f"DEBUG: UnmergeTransform output {j} -> node {node_id}")
769+
else:
770+
stage_output_nodes[output_idx] = node_id
755771

756772
print(f"DEBUG: Creating node {node_id} for stage {stage_idx} transform {i}")
757773

@@ -790,21 +806,35 @@ def build_transformation_graph_from_pytensor(descriptors, variables):
790806
stage_cluster.node(node_id, label, fillcolor="#c0ffc0")
791807
print(f"DEBUG: Added DOT node {node_id} with label {formula_str}")
792808

793-
# Create edges from input nodes to this output node
794-
if isinstance(transform, pytensor.tensor_descriptor.EmbedTransform):
795-
# EmbedTransform: edges from upper_indices (inputs)
796-
for input_idx in upper_indices:
797-
if input_idx in prev_stage_output_nodes:
798-
transform_name = transform.__class__.__name__.replace('Transform', '')
799-
dot.edge(prev_stage_output_nodes[input_idx], node_id,
800-
label=transform_name)
801-
else:
802-
# Other transforms: edges from lower_indices (inputs)
803-
for input_idx in lower_indices:
804-
if input_idx in prev_stage_output_nodes:
805-
transform_name = transform.__class__.__name__.replace('Transform', '')
806-
dot.edge(prev_stage_output_nodes[input_idx], node_id,
807-
label=transform_name)
809+
# Create edges from input nodes to this output node
810+
# Use last_stage_outputs to only connect to the most recent stage
811+
print(f"DEBUG: Creating edges for transform {i} ({transform.__class__.__name__})")
812+
print(f"DEBUG: last_stage_outputs keys: {list(last_stage_outputs.keys())}")
813+
print(f"DEBUG: lower_indices (inputs for non-Embed): {lower_indices}")
814+
print(f"DEBUG: upper_indices (outputs for non-Embed): {upper_indices}")
815+
816+
if isinstance(transform, pytensor.tensor_descriptor.EmbedTransform):
817+
# EmbedTransform: upper_indices are the inputs, lower_indices are the outputs
818+
for input_idx in upper_indices:
819+
if input_idx in last_stage_outputs:
820+
source_node = last_stage_outputs[input_idx]
821+
transform_name = transform.__class__.__name__.replace('Transform', '')
822+
dot.edge(source_node, node_id,
823+
label=transform_name)
824+
print(f"DEBUG: Created edge from {source_node} to {node_id}")
825+
else:
826+
print(f"DEBUG: Input index {input_idx} not found in last_stage_outputs")
827+
else:
828+
# Other transforms: lower_indices are the inputs, upper_indices are the outputs
829+
for input_idx in lower_indices:
830+
if input_idx in last_stage_outputs:
831+
source_node = last_stage_outputs[input_idx]
832+
transform_name = transform.__class__.__name__.replace('Transform', '')
833+
dot.edge(source_node, node_id,
834+
label=transform_name)
835+
print(f"DEBUG: Created edge from {source_node} to {node_id}")
836+
else:
837+
print(f"DEBUG: Input index {input_idx} not found in last_stage_outputs")
808838

809839
except Exception as e:
810840
# Provide more helpful error messages for common issues
@@ -826,38 +856,49 @@ def build_transformation_graph_from_pytensor(descriptors, variables):
826856

827857
for j, output_idx in enumerate(error_output_indices):
828858
node_id = f"s{actual_stage_num}_t{i}_d{output_idx}"
829-
stage_output_nodes[output_idx] = node_id
859+
860+
# For multi-stage examples with UnmergeTransform in stage 0,
861+
# use sequential logical indices to ensure proper mapping to next stage
862+
if (stage_idx == 0 and len(pytensor_descriptors) > 1 and
863+
isinstance(transform, pytensor.tensor_descriptor.UnmergeTransform)):
864+
stage_output_nodes[j] = node_id
865+
else:
866+
stage_output_nodes[output_idx] = node_id
830867
next_formulas[node_id] = sp.Symbol(f"d{output_idx}")
831868
stage_cluster.node(node_id, f"d{output_idx}", fillcolor="#ffcccc")
832869

833870
# Create edges for error case - use same logic as successful case
834871
if isinstance(transform, pytensor.tensor_descriptor.EmbedTransform):
835872
# EmbedTransform: edges from upper_indices (inputs)
836873
for input_idx in upper_indices:
837-
if input_idx in prev_stage_output_nodes:
874+
if input_idx in last_stage_outputs:
838875
transform_name = transform.__class__.__name__.replace('Transform', '')
839-
dot.edge(prev_stage_output_nodes[input_idx], node_id,
876+
dot.edge(last_stage_outputs[input_idx], node_id,
840877
label=transform_name)
841878
else:
842879
# Other transforms: edges from lower_indices (inputs)
843880
for input_idx in lower_indices:
844-
if input_idx in prev_stage_output_nodes:
881+
if input_idx in last_stage_outputs:
845882
transform_name = transform.__class__.__name__.replace('Transform', '')
846-
dot.edge(prev_stage_output_nodes[input_idx], node_id,
883+
dot.edge(last_stage_outputs[input_idx], node_id,
847884
label=transform_name)
848885

849886
# Update for next stage
850887
if stage_output_nodes:
851-
# Merge new stage outputs with existing nodes
888+
# Update prev_stage_output_nodes with all outputs (for final connections)
852889
prev_stage_output_nodes.update(stage_output_nodes)
853890

891+
# Replace last_stage_outputs with only this stage's outputs (for next stage connections)
892+
last_stage_outputs = stage_output_nodes.copy()
893+
854894
# For naive packed descriptors with proper index remapping,
855895
# the UnmergeTransform outputs should now cover all logical dimensions
856896
current_formulas.update(next_formulas)
857-
print(f"DEBUG: Stage {stage_idx} completed, merged stage outputs with existing nodes")
858-
print(f"DEBUG: Available nodes for next stage: {list(prev_stage_output_nodes.keys())}")
897+
print(f"DEBUG: Stage {stage_idx} completed, updated last_stage_outputs with {len(stage_output_nodes)} nodes")
898+
print(f"DEBUG: stage_output_nodes mapping: {stage_output_nodes}")
899+
print(f"DEBUG: Last stage outputs for next stage: {list(last_stage_outputs.keys())}")
900+
859901

860-
861902

862903
# Increment actual stage number only when we actually process a stage
863904
actual_stage_num += 1
@@ -1395,6 +1436,7 @@ def get_length(self, dim):
13951436

13961437
# Process descriptors in REVERSE order for backward graph
13971438
actual_stage_num = len(pytensor_descriptors)
1439+
last_stage_outputs = prev_stage_output_nodes.copy() # Track only the most recent stage's outputs
13981440
for stage_idx in range(len(pytensor_descriptors) - 1, -1, -1):
13991441
tensor_desc = pytensor_descriptors[stage_idx]
14001442
print(f"DEBUG BACKWARD: Processing stage_idx={stage_idx} in reverse")
@@ -1541,7 +1583,17 @@ def get_length(self, dim):
15411583
for j, output_idx in enumerate(backward_output_indices):
15421584
if j < len(output_formulas):
15431585
node_id = f"back_s{actual_stage_num}_t{i}_d{output_idx}"
1544-
stage_output_nodes[output_idx] = node_id
1586+
1587+
# For multi-stage examples, ensure consistent index mapping
1588+
# In backward graph, the last stage (processed first) might need special handling
1589+
if (stage_idx == len(pytensor_descriptors) - 1 and len(pytensor_descriptors) > 1 and
1590+
isinstance(transform, pytensor.tensor_descriptor.MergeTransform)):
1591+
# For MergeTransform outputs in the last stage (first in backward),
1592+
# use sequential indices
1593+
stage_output_nodes[j] = node_id
1594+
print(f"DEBUG BACKWARD: MergeTransform output {j} -> node {node_id}")
1595+
else:
1596+
stage_output_nodes[output_idx] = node_id
15451597

15461598
print(f"DEBUG BACKWARD: Creating node {node_id} for stage {stage_idx} transform {i}")
15471599

@@ -1581,34 +1633,48 @@ def get_length(self, dim):
15811633
print(f"DEBUG BACKWARD: Added DOT node {node_id} with label {formula_str}")
15821634

15831635
# Create edges FROM input nodes TO this output node (normal left-to-right flow)
1636+
# Use last_stage_outputs to only connect to the most recent stage
15841637
for input_idx in backward_input_indices:
1585-
if input_idx in prev_stage_output_nodes:
1638+
logical_idx = hidden_to_logical.get(input_idx, input_idx)
1639+
if logical_idx in last_stage_outputs:
15861640
transform_name = f"{transform.__class__.__name__.replace('Transform', '')}⁻¹"
1587-
dot.edge(prev_stage_output_nodes[input_idx], node_id,
1641+
dot.edge(last_stage_outputs[logical_idx], node_id,
15881642
label=transform_name, color="blue")
15891643

15901644
except Exception as e:
15911645
st.warning(f"Failed to generate backward formula for transform {transform}: {e}")
15921646
# Create error nodes
15931647
for j, output_idx in enumerate(backward_output_indices):
15941648
node_id = f"back_s{actual_stage_num}_t{i}_d{output_idx}"
1595-
stage_output_nodes[output_idx] = node_id
1649+
1650+
# For multi-stage examples, ensure consistent index mapping
1651+
if (stage_idx == len(pytensor_descriptors) - 1 and len(pytensor_descriptors) > 1 and
1652+
isinstance(transform, pytensor.tensor_descriptor.MergeTransform)):
1653+
stage_output_nodes[j] = node_id
1654+
else:
1655+
stage_output_nodes[output_idx] = node_id
15961656
next_formulas[node_id] = sp.Symbol(f"back_d{output_idx}")
15971657
dot.node(node_id, f"back_d{output_idx}", fillcolor="#ffcccc")
15981658

15991659
# Create edges for error case
16001660
for input_idx in backward_input_indices:
1601-
if input_idx in prev_stage_output_nodes:
1661+
logical_idx = hidden_to_logical.get(input_idx, input_idx) if 'hidden_to_logical' in locals() else input_idx
1662+
if logical_idx in last_stage_outputs:
16021663
transform_name = f"{transform.__class__.__name__.replace('Transform', '')}⁻¹"
1603-
dot.edge(prev_stage_output_nodes[input_idx], node_id,
1664+
dot.edge(last_stage_outputs[logical_idx], node_id,
16041665
label=transform_name, color="red")
16051666

16061667
# Update for next stage (going backwards)
16071668
if stage_output_nodes:
1669+
# Update prev_stage_output_nodes with all outputs (for final connections)
16081670
prev_stage_output_nodes.update(stage_output_nodes)
1671+
1672+
# Replace last_stage_outputs with only this stage's outputs (for next stage connections)
1673+
last_stage_outputs = stage_output_nodes.copy()
1674+
16091675
current_formulas.update(next_formulas)
1610-
print(f"DEBUG BACKWARD: Stage {stage_idx} completed, merged stage outputs with existing nodes")
1611-
print(f"DEBUG BACKWARD: Available nodes for next stage: {list(prev_stage_output_nodes.keys())}")
1676+
print(f"DEBUG BACKWARD: Stage {stage_idx} completed, updated last_stage_outputs with {len(stage_output_nodes)} nodes")
1677+
print(f"DEBUG BACKWARD: Last stage outputs for next stage: {list(last_stage_outputs.keys())}")
16121678
actual_stage_num -= 1
16131679
else:
16141680
print(f"DEBUG BACKWARD: Stage {stage_idx} completed with no output nodes")

‎tensor_transforms/examples.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -337,6 +337,22 @@ def get_transform_examples() -> Dict[str, str]:
337337
constexpr auto simple_packed_desc = make_naive_tensor_descriptor_packed(
338338
make_tuple(number<A>{}, number<B>{}, number<C>{}, number<D>{})
339339
);
340+
""",
341+
"Merge Consecutive Dimensions (6D → 3D)": """
342+
constexpr auto packed_6d_desc = make_naive_tensor_descriptor_packed(
343+
make_tuple(number<D0>{}, number<D1>{}, number<D2>{}, number<D3>{}, number<D4>{}, number<D5>{})
344+
);
345+
346+
constexpr auto merged_3d_desc = transform_tensor_descriptor(
347+
packed_6d_desc,
348+
make_tuple(
349+
make_merge_transform(make_tuple(number<D0>{}, number<D1>{})),
350+
make_merge_transform(make_tuple(number<D2>{}, number<D3>{})),
351+
make_merge_transform(make_tuple(number<D4>{}, number<D5>{}))
352+
),
353+
make_tuple(sequence<0, 1>{}, sequence<2, 3>{}, sequence<4, 5>{}),
354+
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{})
355+
);
340356
""",
341357
"Simple Naive Regular (with Strides)": """
342358
constexpr auto simple_strided_desc = make_naive_tensor_descriptor(
@@ -454,6 +470,14 @@ def get_default_variables() -> Dict[str, Dict[str, int]]:
454470
'C': 4,
455471
'D': 5,
456472
},
473+
"Merge Consecutive Dimensions (6D → 3D)": {
474+
'D0': 2,
475+
'D1': 3,
476+
'D2': 4,
477+
'D3': 5,
478+
'D4': 6,
479+
'D5': 7,
480+
},
457481
"Simple Naive Regular (with Strides)": {
458482
'M': 4,
459483
'N': 8,

0 commit comments

Comments
 (0)