nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

update doc example for multinomial (#10657)

Browse Source 
also added many `s` for consistency
This commit is contained in:
chenyu
2025-06-05 20:16:52 -04:00
committed by GitHub
parent 61352b8aa2
commit e67642d430
1 changed files with 21 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
+21 -20
View File
@@ -446,7 +446,7 @@ class Tensor(MathTrait):
@staticmethod
def from_url(url:str, gunzip:bool=False, **kwargs) -> Tensor:
"""
Create a Tensor from a URL.
Creates a Tensor from a URL.
This is the preferred way to access Internet resources.
It currently returns a DISK Tensor, but in the future it may return an HTTP Tensor.
@@ -869,11 +869,11 @@ class Tensor(MathTrait):
@staticmethod
def randperm(n:int, device=None, dtype=dtypes.int32, **kwargs) -> Tensor:
"""
Return a tensor with a random permutation of integers from 0 to n-1.
Returns a tensor with a random permutation of integers from `0` to `n-1`.
```python exec="true" source="above" session="tensor" result="python"
Tensor.manual_seed(42)
print(Tensor.randperm(4).numpy())
print(Tensor.randperm(6).numpy())
```
"""
r = Tensor.rand(n, device=device, **kwargs)
@@ -882,12 +882,13 @@ class Tensor(MathTrait):
def multinomial(self:Tensor, num_samples:int = 1, replacement:bool = False) -> Tensor:
"""
Sample from a multinomial distribution weighted by `self`.
Returns a tensor with `num_samples` indices sampled from a multinomial distribution weighted by `self`.
NOTE: `replacement=False` for `num_samples > 1` is not supported yet.
```python exec="true" source="above" session="tensor" result="python"
Tensor.manual_seed(42)
t = Tensor.arange(10)
print(t.multinomial().numpy())
t = Tensor([1, 2, 3, 4])
print(t.multinomial(20, replacement=True).numpy())
```
"""
assert 1 <= self.ndim <= 2 and num_samples > 0, f"{self.ndim=} must be 1 or 2 dim, {num_samples=} must be positive"
@@ -902,7 +903,7 @@ class Tensor(MathTrait):
def gradient(self, *targets:Tensor, gradient:Tensor|None=None, materialize_grads=False) -> list[Tensor]:
"""
Compute the gradient of the targets with respect to self.
Computes the gradient of the targets with respect to self.
```python exec="true" source="above" session="tensor" result="python"
x = Tensor.eye(3)
@@ -1204,7 +1205,7 @@ class Tensor(MathTrait):
def __getitem__(self, indices) -> Tensor:
"""
Retrieve a sub-tensor using indexing.
Retrieves a sub-tensor using indexing.
Supported Index Types: `int | slice | Tensor | None | list | tuple | Ellipsis`
@@ -1316,7 +1317,7 @@ class Tensor(MathTrait):
def repeat_interleave(self, repeats:int, dim:int|None=None) -> Tensor:
"""
Repeat elements of a tensor.
Repeats elements of a tensor.
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([1, 2, 3])
@@ -1625,7 +1626,7 @@ class Tensor(MathTrait):
def masked_fill(self:Tensor, mask:Tensor, value:Tensor|ConstType) -> Tensor:
"""
Replace `self` with `value` wherever the elements of `mask` are True.
Replaces `self` with `value` wherever the elements of `mask` are True.
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([1, 2, 3, 4, 5])
@@ -2807,7 +2808,7 @@ class Tensor(MathTrait):
def fuse(self) -> Tensor:
"""
Make this a single kernel back to Ops.CONTIGUOUS on the inputs.
Makes this a single kernel back to Ops.CONTIGUOUS on the inputs.
Useful for single kernel softmax and flash attention.
Careful, this can break codegen or make kernels really slow.
@@ -3127,7 +3128,7 @@ class Tensor(MathTrait):
def reciprocal(self) -> Tensor:
"""
Compute `1/x` element-wise.
Computes `1/x` element-wise.
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([1., 2., 3., 4.]).reciprocal().numpy())
@@ -3571,7 +3572,7 @@ class Tensor(MathTrait):
def bitwise_and(self, x:Tensor|ConstType, reverse=False) -> Tensor:
"""
Compute the bitwise AND of `self` and `x`.
Computes the bitwise AND of `self` and `x`.
Equivalent to `self & x`.
Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
@@ -3586,7 +3587,7 @@ class Tensor(MathTrait):
def bitwise_or(self, x:Tensor|ConstType, reverse=False) -> Tensor:
"""
Compute the bitwise OR of `self` and `x`.
Computes the bitwise OR of `self` and `x`.
Equivalent to `self | x`.
Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
@@ -3601,7 +3602,7 @@ class Tensor(MathTrait):
def bitwise_not(self) -> Tensor:
"""
Compute the bitwise NOT of `self`.
Computes the bitwise NOT of `self`.
Equivalent to `~self`.
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([0, 2, 5, 255], dtype="int8").bitwise_not().numpy())
@@ -3689,7 +3690,7 @@ class Tensor(MathTrait):
def where(self:Tensor, x:Tensor|ConstType|sint, y:Tensor|ConstType|sint) -> Tensor:
"""
Return a tensor of elements selected from either `x` or `y`, depending on `self`.
Returns a tensor of elements selected from either `x` or `y`, depending on `self`.
`output_i = x_i if self_i else y_i`.
```python exec="true" source="above" session="tensor" result="python"
@@ -3713,7 +3714,7 @@ class Tensor(MathTrait):
def copysign(self, other) -> Tensor:
"""
Return a tensor of with the magnitude of `self` and the sign of `other`, elementwise.
Returns a tensor of with the magnitude of `self` and the sign of `other`, elementwise.
"""
# NOTE: torch always return in float, we return based on the broadcasting rule.
other = self._broadcasted(other)[1]
@@ -3952,7 +3953,7 @@ class Tensor(MathTrait):
def cross_entropy(self, Y:Tensor, reduction:ReductionStr="mean", label_smoothing:float=0.0) -> Tensor:
"""
Compute the cross entropy loss between input logits and target.
Computes the cross entropy loss between input logits and target.
NOTE: `self` are logits and `Y` are the target labels or class probabilities.
@@ -3977,7 +3978,7 @@ class Tensor(MathTrait):
def nll_loss(self, Y:Tensor, weight:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean") -> Tensor:
"""
Compute the negative log likelihood loss between log-probabilities and target labels.
Computes the negative log likelihood loss between log-probabilities and target labels.
NOTE: `self` is log-probabilities and `Y` is the Y labels or class probabilities.
@@ -4060,7 +4061,7 @@ class Tensor(MathTrait):
def size(self, dim:int|None=None) -> sint|tuple[sint, ...]:
"""
Return the size of the tensor. If `dim` is specified, return the length along dimension `dim`. Otherwise return the shape of the tensor.
Returns the size of the tensor. If `dim` is specified, return the length along dimension `dim`. Otherwise return the shape of the tensor.
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([[4, 5, 6], [7, 8, 9]])