issue/288: improve the compatibility of the torch implementations of gemm and random sample

test/infiniop/gemm.py

@@ -57,11 +57,14 @@ infiniopGemmDescriptor_t = POINTER(GemmDescriptor)
 
 # 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)