code: def test_expand_trace_observations(self):
      for batch_shape, dtype in itertools.product(
          ([], [2]), (torch.float, torch.double)
      ):
          q, d = 3, 4
          X = torch.rand(
              *batch_shape, q, d, device=self.device, dtype=dtype
          )
          # test nullop behavior
          self.assertTrue(torch.equal(expand_trace_observations(X), X))
          self.assertTrue(
              torch.equal(
                  expand_trace_observations(X, fidelity_dims=[1]), X
              )
          )
          # test default behavior
          num_tr = 2
          X_expanded = expand_trace_observations(
              X, num_trace_obs=num_tr
          )
          self.assertEqual(
              X_expanded.shape,
              torch.Size(batch_shape + [q * (1 + num_tr), d]),
          )
          for i in range(num_tr):
              X_sub = X_expanded[..., q * i : q * (i + 1), :]
              self.assertTrue(torch.equal(X_sub[..., :-1], X[..., :-1]))
              X_sub_expected = (1 - i / (num_tr + 1)) * X[..., :q, -1]
              self.assertTrue(
                  torch.equal(X_sub[..., -1], X_sub_expected)
              )
          # test custom fidelity dims
          fdims = [0, 2]
          num_tr = 3
          X_expanded = expand_trace_observations(
              X, fidelity_dims=fdims, num_trace_obs=num_tr
          )
          self.assertEqual(
              X_expanded.shape,
              torch.Size(batch_shape + [q * (1 + num_tr), d]),
          )
          for j, i in itertools.product([1, 3], range(num_tr)):
              X_sub = X_expanded[..., q * i : q * (i + 1), j]
              self.assertTrue(torch.equal(X_sub, X[..., j]))
          for j, i in itertools.product(fdims, range(num_tr)):
              X_sub = X_expanded[..., q * i : q * (i + 1), j]
              X_sub_expected = (1 - i / (1 + num_tr)) * X[..., :q, j]
              self.assertTrue(torch.equal(X_sub, X_sub_expected))
          # test gradients
          num_tr = 2
          fdims = [1]
          X.requires_grad_(True)
          X_expanded = expand_trace_observations(
              X, fidelity_dims=fdims, num_trace_obs=num_tr
          )
          out = X_expanded.sum()
          out.backward()
          grad_exp = torch.full_like(X, 1 + num_tr)
          grad_exp[..., fdims] = 1 + sum(
              (i + 1) / (num_tr + 1) for i in range(num_tr)
          )
          self.assertAllClose(X.grad, grad_exp)
  
imports: import itertools
  from unittest import mock
  from unittest.mock import patch
  import torch
  from botorch.acquisition.objective import GenericMCObjective, LearnedObjective
  from botorch.acquisition.utils import compute_best_feasible_objective, expand_trace_observations, get_acquisition_function, get_infeasible_cost, get_optimal_samples, project_to_sample_points, project_to_target_fidelity, prune_inferior_points, repeat_to_match_aug_dim
  from botorch.exceptions.errors import BotorchTensorDimensionError, DeprecationError, UnsupportedError
  from botorch.models import SingleTaskGP
  from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
  from gpytorch.distributions import MultivariateNormal
label_desc: Use image node_id 2 for calling the TestFidelityUtils obj's underlying member method code with example usage: obj.test_expand_trace_observations() without return types