code: def test_shape(device, levy_area_approximation, return_U, return_A):
      if device == gpu and not torch.cuda.is_available():
          pytest.skip(msg="CUDA not available.")
  
      for shape, A_shape in (
          ((SMALL_BATCH_SIZE, D), (SMALL_BATCH_SIZE, D, D)),
          ((SMALL_BATCH_SIZE,), (SMALL_BATCH_SIZE,)),
          ((), ()),
      ):
          ta, tb, bm = _setup(device, levy_area_approximation, shape)
          sample1 = bm(ta, return_U=return_U, return_A=return_A)
          sample2 = bm(tb, return_U=return_U, return_A=return_A)
          sample3 = bm(ta, tb, return_U=return_U, return_A=return_A)
          shapes = []
          A_shapes = []
          for sample in (sample1, sample2, sample3):
              if return_U:
                  if return_A:
                      W1, U1, A1 = sample
                      shapes.append(W1.shape)
                      shapes.append(U1.shape)
                      A_shapes.append(A1.shape)
                  else:
                      W1, U1 = sample
                      shapes.append(W1.shape)
                      shapes.append(U1.shape)
              else:
                  if return_A:
                      W1, A1 = sample
                      shapes.append(W1.shape)
                      A_shapes.append(A1.shape)
                  else:
                      W1 = sample
                      shapes.append(W1.shape)
  
          for shape_ in shapes:
              assert shape_ == shape
          for shape_ in A_shapes:
              assert shape_ == A_shape
  
file_path: torchsde/tests/test_brownian_interval.py
imports: import sys
  import math
  import numpy.random
  import torch
  from scipy.stats import kstest
  import pytest
  import torchsde
label_desc: Use image node_id 4 for calling a global function with example usage: test_shape(device, levy_area_approximation, return_U, return_A) without return types