add imports

scikit-hep · Sep 17, 2024 · 70a2630 · 70a2630
1 parent 14b09d2
commit 70a2630
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/src/awkward/operations/ak_to_jaggedtensor.py b/src/awkward/operations/ak_to_jaggedtensor.py
@@ -4,9 +4,14 @@
 
 import awkward as ak
 from awkward._dispatch import high_level_function
+from awkward._nplikes.numpy import Numpy
+from awkward._nplikes.numpy_like import NumpyMetadata
 
 __all__ = ("to_jaggedtensor",)
 
+numpy = Numpy.instance()
+np = NumpyMetadata.instance()
+
 
 @high_level_function()
 def to_jaggedtensor(array, padded=False, padding_value=0, max_lengths=None, keep_regular=True):

diff --git a/tests/test_3246_to_jaggedtensor_from_jaggedtensor.py b/tests/test_3246_to_jaggedtensor_from_jaggedtensor.py
@@ -9,8 +9,8 @@
 
 to_jaggedtensor = ak.operations.to_jaggedtensor
 
-tf = pytest.importorskip("torch")
-tf = pytest.importorskip("fbgemm_gpu")
+torch = pytest.importorskip("torch")
+fbgemm_gpu = pytest.importorskip("fbgemm_gpu")
 
 content = ak.contents.NumpyArray(
     np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])