Internal

PiperOrigin-RevId: 948049485

Author: Pathways-on-Cloud Team

pathwaysutils/profiling.py

@@ -64,25 +64,8 @@ def call_profile_executable(self) -> None:
       raise RuntimeError(
           "_call_profile_executable called with no active executable."
       )
-    # If the profile request contains xprofTraceOptions, then we need to pass
-    # out_avals and out_shardings to the executable call because the
-    # executable will return a future that needs to be resolved. This is true
-    # for both starting and stopping a trace.
-    if (
-        self.profile_request is not None
-        and "xprofTraceOptions" in self.profile_request
-    ):
-      out_avals = [jax.core.ShapedArray((1,), jnp.object_)]
-      out_shardings = [
-          getattr(
-              jax.sharding,
-              "make_single_device_sharding",
-              jax.sharding.SingleDeviceSharding,
-          )(backend.get_default_device())
-      ]
-    else:
-      out_avals = ()
-      out_shardings = ()
+    out_avals = ()
+    out_shardings = ()
 
     _, result_future = self.executable.call(
         out_avals=out_avals, out_shardings=out_shardings

pathwaysutils/test/profiling_test.py

@@ -389,10 +389,8 @@ def test_stop_trace_with_xprof_options_passes_out_avals(self):
 
     with self.subTest("out_avals_properties"):
       _, kwargs = self.mock_plugin_executable_cls.return_value.call.call_args
-      self.assertLen(kwargs["out_avals"], 1)
-      (out_aval,) = kwargs["out_avals"]
-      self.assertEqual(out_aval.shape, (1,))
-      self.assertEqual(out_aval.dtype, jnp.object_)
+      self.assertEqual(kwargs["out_avals"], ())
+      self.assertEqual(kwargs["out_shardings"], ())
 
   def test_stop_trace_before_start_error(self):
     with self.assertRaisesRegex(