code: def test_gmm_params_energy():
      phi, mu, cov, L, log_det_cov = gmm_params(z, gamma)
      assert phi.numpy().shape[0] == K == log_det_cov.shape[0]
      assert mu.numpy().shape == (K, D)
      assert cov.numpy().shape == L.numpy().shape == (K, D, D)
      for _ in range(cov.numpy().shape[0]):
          assert (np.diag(cov[_].numpy()) >= 0.0).all()
          assert (np.diag(L[_].numpy()) >= 0.0).all()
  
      sample_energy, cov_diag = gmm_energy(
          z, phi, mu, cov, L, log_det_cov, return_mean=True
      )
      assert sample_energy.numpy().shape == cov_diag.numpy().shape == ()
  
      sample_energy, _ = gmm_energy(
          z, phi, mu, cov, L, log_det_cov, return_mean=False
      )
      assert sample_energy.numpy().shape[0] == N
  
label_desc: Use image node_id 1 for calling a global function with example usage: test_gmm_params_energy() without return types