idea does not work

andrinr · andrinr · commit 500fae6d2ebb · 2025-11-04T16:31:22.000+01:00
diff --git a/run_12f4be7b-ff64-48b9-81bf-e74dd93bc282/logs/metrics.csv b/run_12f4be7b-ff64-48b9-81bf-e74dd93bc282/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_39ca8eca-c4ec-4e7b-b3dd-ad520d1f4801/logs/metrics.csv b/run_39ca8eca-c4ec-4e7b-b3dd-ad520d1f4801/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_3cca4752-436d-4eca-b57a-ac93424b252b/logs/metrics.csv b/run_3cca4752-436d-4eca-b57a-ac93424b252b/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_421ea214-37f2-4294-8f41-a930506e2eee/logs/metrics.csv b/run_421ea214-37f2-4294-8f41-a930506e2eee/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_563e8a0c-5f99-461d-97b6-e8b7f12eaf2e/logs/metrics.csv b/run_563e8a0c-5f99-461d-97b6-e8b7f12eaf2e/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_5e957f9e-ff75-4384-ab35-883ae4d6378f/logs/metrics.csv b/run_5e957f9e-ff75-4384-ab35-883ae4d6378f/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_614372b4-296f-4d40-8d1b-e59113663edd/logs/metrics.csv b/run_614372b4-296f-4d40-8d1b-e59113663edd/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_62b9044c-a83f-416f-ba0c-cfcf3ab0bf01/logs/metrics.csv b/run_62b9044c-a83f-416f-ba0c-cfcf3ab0bf01/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_66bd8d02-acd1-48ff-9f11-c58e7b8c5621/logs/metrics.csv b/run_66bd8d02-acd1-48ff-9f11-c58e7b8c5621/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_68e91784-1fb9-4ea7-b268-deacda9c2d8b/logs/metrics.csv b/run_68e91784-1fb9-4ea7-b268-deacda9c2d8b/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_6eea3bf7-6c96-4e47-911e-56b368deeb86/logs/metrics.csv b/run_6eea3bf7-6c96-4e47-911e-56b368deeb86/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_7b823032-5ab2-4a84-822a-55f943db1455/logs/metrics.csv b/run_7b823032-5ab2-4a84-822a-55f943db1455/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_7fafb835-fa09-4fb9-ac62-f33e9d9e7405/logs/metrics.csv b/run_7fafb835-fa09-4fb9-ac62-f33e9d9e7405/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_8bbb74e8-1355-40f0-89cd-c49d7f8c78af/logs/metrics.csv b/run_8bbb74e8-1355-40f0-89cd-c49d7f8c78af/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_8d6bacf4-8cde-4b30-886b-d14ef10c7d6c/logs/metrics.csv b/run_8d6bacf4-8cde-4b30-886b-d14ef10c7d6c/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_9466e74a-6dcc-4b4b-8cfc-b06fbaf912c9/logs/metrics.csv b/run_9466e74a-6dcc-4b4b-8cfc-b06fbaf912c9/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_9de6bb4b-9e0d-44f9-94e0-754aeb461d56/logs/metrics.csv b/run_9de6bb4b-9e0d-44f9-94e0-754aeb461d56/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_a71217f8-3fe7-44a0-a8e2-350dde160320/logs/metrics.csv b/run_a71217f8-3fe7-44a0-a8e2-350dde160320/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_abe9ca59-56a0-4857-9a43-176f0064c8f9/logs/metrics.csv b/run_abe9ca59-56a0-4857-9a43-176f0064c8f9/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_af0f872b-0361-47be-961c-e2858c6f0db4/logs/metrics.csv b/run_af0f872b-0361-47be-961c-e2858c6f0db4/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_c4baea3d-588b-497c-8c44-5cbedd9b64f6/logs/metrics.csv b/run_c4baea3d-588b-497c-8c44-5cbedd9b64f6/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_c9eb768a-3e05-4030-995e-4d45ddb26519/logs/metrics.csv b/run_c9eb768a-3e05-4030-995e-4d45ddb26519/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_e885dda5-70b5-468a-909f-dc7bcbff8e38/logs/metrics.csv b/run_e885dda5-70b5-468a-909f-dc7bcbff8e38/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/run_f602a2cb-97cd-40b7-9a15-9d7a762482da/logs/metrics.csv b/run_f602a2cb-97cd-40b7-9a15-9d7a762482da/logs/metrics.csv
@@ -0,0 +1 @@
+timestamp,key,value,step
diff --git a/tesseract_jax/primitive.py b/tesseract_jax/primitive.py
@@ -332,8 +332,7 @@ def _unpack_hashable(obj: _Hashable) -> Any:
 
 
 def apply_tesseract(
-    tesseract_client: Tesseract,
-    inputs: Any,
+    tesseract_client: Tesseract, inputs: Any, static_arg_names: Sequence[str] = ()
 ) -> Any:
     """Applies the given Tesseract object to the inputs.
 
@@ -367,6 +366,9 @@ def apply_tesseract(
     Args:
         tesseract_client: The Tesseract object to apply.
         inputs: The inputs to apply to the Tesseract object.
+        static_arg_names: Names of input arguments that should be treated as static
+            (i.e., not traced). This is useful for arguments that affect control flow
+            or shape but are not arrays.
 
     Returns:
         The outputs of the Tesseract object after applying the inputs.
@@ -401,7 +403,18 @@ def apply_tesseract(
     client = Jaxeract(tesseract_client)
 
     flat_args, input_pytreedef = jax.tree.flatten(inputs)
-    is_static_mask = tuple(_is_static(arg) for arg in flat_args)
+    is_static_mask = tuple((_is_static(arg)) for arg in flat_args)
+    paths_and_values, _ = jax.tree.flatten_with_path(inputs)
+    paths_concatenated = [
+        "/".join(key.key for key in path) for path, _ in paths_and_values
+    ]
+    is_static_mask_ = tuple(path in static_arg_names for path in paths_concatenated)
+
+    # or with the previous mask
+    is_static_mask = tuple(
+        a or b for a, b in zip(is_static_mask, is_static_mask_, strict=True)
+    )
+
     array_args, static_args = split_args(flat_args, is_static_mask)
     static_args = tuple(_make_hashable(arg) for arg in static_args)
 
diff --git a/tests/test_endtoend.py b/tests/test_endtoend.py
@@ -582,13 +582,15 @@ def loss_fn(a):
         )
 
         c = vectoradd_fn_a(a)["c"]
+        c = jax.lax.stop_gradient(c)
 
         vectoradd_fn_b: jax.Callable = lambda a: apply_tesseract(
             vectoradd_tess,
             inputs=dict(
                 a=a,
                 b=c,
             ),
+            static_arg_names=["b"],
         )
 
         outputs = vectoradd_fn_b(a)
