code: def test_input_across_dims(
      self,
      N: int,
      T: int,
      Q: int,
      R: int,
      V: int,
      D: int,
      dim: int,
      batch_first: bool,
  ):
      if dim == 1:  # TODO: fix when dim is 1
          size = [T]
      elif dim == 2:
          size = [N, T]
      elif dim == 3:
          size = [N, T, Q]
      elif dim == 4:
          size = [N, T, Q, R]
  
      emb = nn.Embedding(V, D)
      x = torch.randint(low=0, high=V, size=size)
      self.run_test(
          x, emb, batch_first=batch_first, ew_compatible=N > 0
      )
  
file_path: opacus/opacus/tests/grad_samples/embedding_test.py
imports: import hypothesis.strategies
  import torch
  import torch.nn
  from hypothesis import given, settings
  from .common import GradSampleHooks_test
label_desc: Use image node_id 1 for calling the Embedding_test obj's underlying member method code with example usage: obj.test_input_across_dims(N, T, Q, R, V, D, dim, batch_first) without return types