Merge pull request #289 from InfiniTensor/issue/288_improve_torch_implementation_compatibility

issue/288: Improve the Compatibility of the Torch Implementations

Author: PanZezhong1725

test/infiniop/gemm.py

@@ -57,11 +57,14 @@ class GemmDescriptor(Structure):
 
 # PyTorch implementation for matrix multiplication
 def gemm(d, _c, beta, _a, _b, alpha):
-    if _c.ndim == 2:
-        torch.addmm(_c, _a, _b, beta=beta, alpha=alpha, out=d)
-    elif _c.ndim == 3:
-        torch.baddbmm(_c, _a, _b, beta=beta, alpha=alpha, out=d)
-    else:
+    try:
+        if _c.ndim == 2:
+            torch.addmm(_c, _a, _b, beta=beta, alpha=alpha, out=d)
+        elif _c.ndim == 3:
+            torch.baddbmm(_c, _a, _b, beta=beta, alpha=alpha, out=d)
+        else:
+            raise
+    except Exception:
         torch.matmul(_a, _b, out=d)
         d.mul_(alpha).add_(_c, alpha=beta)
 

test/infiniop/random_sample.py

@@ -67,8 +67,13 @@ def random_sample(data, random_val, topp, topk, voc, temperature):
 
         k_index = min(topk, voc) - 1
         threshold = min(cum_probs[k_index], topp) * random_val
-
-        idx = torch.searchsorted(cum_probs, threshold)
+        
+        try:
+            idx = torch.searchsorted(cum_probs, threshold)
+        except Exception:
+            # Fallback for manual search if torch.searchsorted is not supported
+            indices = (cum_probs >= threshold).nonzero(as_tuple=True)[0]
+            idx = indices[0] if indices.numel() > 0 else torch.tensor(len(cum_probs)-1, device=cum_probs.device)
         return sorted_indices[idx]
 
     return torch.argmax(data)

test/infiniop/rearrange.py

@@ -116,11 +116,11 @@ def column_major_strides(shape):
 NUM_ITERATIONS = 1000
 
 
-class RerrangeDescriptor(Structure):
+class RearrangeDescriptor(Structure):
     _fields_ = [("device", c_int32)]
 
 
-infiniopRearrangeDescriptor_t = POINTER(RerrangeDescriptor)
+infiniopRearrangeDescriptor_t = POINTER(RearrangeDescriptor)
 
 
 def rearrange_torch(x, x_shape, y_stride):