code: def assert_actor_shapes_are_correct(
      self, actor_output: continuous_networks.SacActorOutput
  ):
      with self.subTest(name="mean_action"):
          self.assertEqual(
              actor_output.mean_action.shape, self.action_shape
          )
      with self.subTest(name="sampled_action"):
          self.assertEqual(
              actor_output.sampled_action.shape, self.action_shape
          )
      with self.subTest(name="log_probability"):
          self.assertEqual(actor_output.log_probability.shape, ())
  
file_path: dopamine/tests/dopamine/jax/continuous_networks_test.py
imports: from absl.testing import absltest
  from dopamine.jax import continuous_networks
  import jax
  from jax import numpy
label_desc: Use image node_id 2 for calling the SacNetworkTest obj's underlying member method code with example usage: obj.assert_actor_shapes_are_correct(actor_output) without return types