docstr: def _test_against_numpy(t, indices, dim):
      actual = torch.take_along_dim(t, indices, dim=dim)
      t_np = t.cpu().numpy()
      indices_np = indices.cpu().numpy()
      expected = np.take_along_axis(t_np, indices_np, axis=dim)
      numpy_testing_assert_equal_helper(actual, expected)
docstr_tok: ["def","_test_against_numpy","(","t",",","indices",",","dim",")",":","actual","=","torch.take_along_dim","(","t",",","indices",",","dim=dim",")","t_np","=","t.cpu","(",")",".numpy","(",")","indices_np","=","indices.cpu","(",")",".numpy","(",")","expected","=","np.take_along_axis","(","t_np",",","indices_np",",","axis=dim",")","numpy_testing_assert_equal_helper","(","actual",",","expected",")"]
code: def test_take_along_dim(self):
      """
      def _test_against_numpy(t, indices, dim):
          actual = torch.take_along_dim(t, indices, dim=dim)
          t_np = t.cpu().numpy()
          indices_np = indices.cpu().numpy()
          expected = np.take_along_axis(t_np, indices_np, axis=dim)
          numpy_testing_assert_equal_helper(actual, expected)
      """
  
      def _test_against_numpy(t: Tensor, indices: Tensor, dim):
          actual = t.gather(indices, dim=dim)
          t_np = t.numpy()
          indices_np = indices.numpy()
          expected = np.take_along_axis(t_np, indices_np, axis=dim)
          numpy_testing_assert_equal_helper(actual, expected)
  
          # TODO argsort
          """
        for shape in [(3, 2), (2, 3, 5), (2, 4, 0), (2, 3, 1, 4)]:
          for noncontiguous in [True, False]:
            for dtype in (dtypes.float32, dtypes.int64):
              t = make_tensor(shape, dtype=dtype, noncontiguous=noncontiguous)
              for dim in list(range(t.ndim)) + [None]:
                if dim is None:
                  indices = argsort(t.reshape(-1))
                else:
                  indices = argsort(t, dim=dim)
  
            _test_against_numpy(t, indices, dim)
        """
  
          # test broadcasting
          t = Tensor.ones((3, 4, 1))
          indices = Tensor.ones((1, 2, 5), dtype=dtypes.int64)
  
          _test_against_numpy(t, indices, 1)
  
          # test empty indices
          t = Tensor.ones((3, 4, 5))
          indices = Tensor.ones((3, 0, 5), dtype=dtypes.int64)
  
          _test_against_numpy(t, indices, 1)
  
imports: import math, unittest, random, copy
  import numpy
  from tinygrad import Tensor, dtypes, Device
  from tinygrad.lazy import LazyBuffer
  from tinygrad.shape.shapetracker import ShapeTracker
  from tinygrad.shape.view import View
  from tinygrad.helpers import CI
label_desc: Use image node_id 21 for calling the TestIndexing obj's underlying member method code with example usage: obj.test_take_along_dim() without return types