just get dtype from kwargs (#4355)

2024-04-30 16:26:14 +09:00
parent a9a1fa6bbf
commit a2d81514fd
1 changed files with 2 additions and 1 deletions
@@ -514,7 +514,8 @@ class Tensor:
    t = Tensor.randint(2, 3, low=5, high=10)
    print(t.numpy())
    """
-    return Tensor.uniform(*shape, low=low, high=high, dtype=dtypes.int32, **kwargs)
+    assert dtypes.is_int(dtype := kwargs.pop("dtype", dtypes.int32)), f"Unsupported dtype {dtype} for randint"
+    return Tensor.uniform(*shape, low=low, high=high, dtype=dtype, **kwargs)

  @staticmethod
  def normal(*shape, mean=0.0, std=1.0, **kwargs) -> Tensor: