update tree_map API call (GoogleCloudPlatform#838)

A9isha · web-flow · commit 7a8fab2cf45f · 2025-08-19T18:18:57.000-07:00
diff --git a/dags/multipod/legacy_tests/gpt1-like.py b/dags/multipod/legacy_tests/gpt1-like.py
@@ -153,7 +153,7 @@ def calculate_tflops(f, *args, **kwargs):
 
 def training_loop(in_act, in_layers):
   _, grad_layers = multiply_layers_and_grad(in_act, in_layers)
-  out_layers = jax.tree_map(
+  out_layers = jax.tree_util.tree_map(
       lambda param, grad: param - 1e-4 * grad, in_layers, grad_layers[0]
   )
   return out_layers

Original file line number	Diff line number	Diff line change
`@@ -153,7 +153,7 @@ def calculate_tflops(f, args, *kwargs):`
`153`	`153`
`154`	`154`	`def training_loop(in_act, in_layers):`
`155`	`155`	`_, grad_layers = multiply_layers_and_grad(in_act, in_layers)`
`156`		`- out_layers = jax.tree_map(`
	`156`	`+ out_layers = jax.tree_util.tree_map(`
`157`	`157`	`lambda param, grad: param - 1e-4 * grad, in_layers, grad_layers[0]`
`158`	`158`	`)`
`159`	`159`	`return out_layers`