diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index 2a39b85..20b9216 100644 --- a/pathwaysutils/profiling.py +++ b/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 diff --git a/pathwaysutils/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index 0909919..aad79e6 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/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(