Complex Ops
Reduce¤
sum
¤
sum(
axis: Optional[Union[int, Sequence[int]]] = None,
keepdim=False,
acc_dtype: Optional[DTypeLike] = None,
)
Returns the sum of the elements of the tensor along the specified axis or axes.
You can pass in axis
and keepdim
keyword arguments to control the axis along
which the maximum is computed and whether the reduced dimensions are retained.
You can pass in acc_dtype
keyword argument to control the data type of the accumulation.
If not specified, the accumulation data type is chosen based on the input tensor's data type.
t = Tensor.arange(6).reshape(2, 3)
print(t.numpy())
[[0 1 2]
[3 4 5]]
print(t.sum().numpy())
15
print(t.sum(axis=0).numpy())
[3 5 7]
print(t.sum(axis=1).numpy())
[ 3 12]
Source code in tinygrad/tensor.py
1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469 1470 1471 |
|
prod
¤
prod(
axis: Optional[Union[int, Sequence[int]]] = None,
keepdim=False,
acc_dtype: Optional[DTypeLike] = None,
)
Returns the product of the elements of the tensor along the specified axis or axes.
You can pass in axis
and keepdim
keyword arguments to control the axis along
which the maximum is computed and whether the reduced dimensions are retained.
You can pass in acc_dtype
keyword argument to control the data type of the accumulation.
If not specified, the accumulation data type is chosen based on the input tensor's data type.
t = Tensor([-1, -2, -3, 1, 2, 3]).reshape(2, 3)
print(t.numpy())
[[-1 -2 -3]
[ 1 2 3]]
print(t.prod().numpy())
-36
print(t.prod(axis=0).numpy())
[-1 -4 -9]
print(t.prod(axis=1).numpy())
[-6 6]
Source code in tinygrad/tensor.py
1473 1474 1475 1476 1477 1478 1479 1480 1481 1482 1483 1484 1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 |
|
max
¤
Returns the maximum value of the tensor along the specified axis or axes.
You can pass in axis
and keepdim
keyword arguments to control the axis along
which the maximum is computed and whether the reduced dimensions are retained.
t = Tensor([[1, 0, 2], [5, 4, 3]])
print(t.numpy())
[[1 0 2]
[5 4 3]]
print(t.max().numpy())
5
print(t.max(axis=0).numpy())
[5 4 3]
print(t.max(axis=1, keepdim=True).numpy())
[[2]
[5]]
Source code in tinygrad/tensor.py
1499 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 |
|
min
¤
Returns the minimum value of the tensor along the specified axis or axes.
You can pass in axis
and keepdim
keyword arguments to control the axis along
which the minimum is computed and whether the reduced dimensions are retained.
t = Tensor([[1, 0, 2], [5, 4, 3]])
print(t.numpy())
[[1 0 2]
[5 4 3]]
print(t.min().numpy())
0
print(t.min(axis=0).numpy())
[1 0 2]
print(t.min(axis=1, keepdim=True).numpy())
[[0]
[3]]
Source code in tinygrad/tensor.py
1522 1523 1524 1525 1526 1527 1528 1529 1530 1531 1532 1533 1534 1535 1536 1537 1538 1539 1540 1541 1542 1543 |
|
any
¤
Tests if any element evaluates to True
along the specified axis or axes.
You can pass in axis
and keepdim
keyword arguments to control the reduce axis and whether the reduced dimensions are retained.
t = Tensor([[True, True], [True, False], [False, False]])
print(t.numpy())
[[ True True]
[ True False]
[False False]]
print(t.any().numpy())
True
print(t.any(axis=0).numpy())
[ True True]
print(t.any(axis=1, keepdim=True).numpy())
[[ True]
[ True]
[False]]
Source code in tinygrad/tensor.py
1545 1546 1547 1548 1549 1550 1551 1552 1553 1554 1555 1556 1557 1558 1559 1560 1561 1562 1563 1564 1565 |
|
all
¤
Tests if all element evaluates to True
along the specified axis or axes.
You can pass in axis
and keepdim
keyword arguments to control the reduce axis and whether the reduced dimensions are retained.
t = Tensor([[True, True], [True, False], [False, False]])
print(t.numpy())
[[ True True]
[ True False]
[False False]]
print(t.all().numpy())
False
print(t.all(axis=0).numpy())
[False False]
print(t.all(axis=1, keepdim=True).numpy())
[[ True]
[False]
[False]]
Source code in tinygrad/tensor.py
1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585 1586 1587 |
|
mean
¤
Returns the mean value of the tensor along the specified axis or axes.
You can pass in axis
and keepdim
keyword arguments to control the axis along
which the mean is computed and whether the reduced dimensions are retained.
Tensor.manual_seed(42)
t = Tensor.normal(2, 3, mean=2.5, std=0.5)
print(t.numpy())
[[2.9889 2.7339 2.7763]
[2.3356 2.0722 2.6376]]
print(t.mean().numpy())
2.5907671
print(t.mean(axis=0).numpy())
[2.6623 2.4031 2.707 ]
print(t.mean(axis=1).numpy())
[2.833 2.3485]
Source code in tinygrad/tensor.py
1589 1590 1591 1592 1593 1594 1595 1596 1597 1598 1599 1600 1601 1602 1603 1604 1605 1606 1607 1608 1609 1610 1611 1612 1613 |
|
var
¤
Returns the variance of the tensor along the specified axis or axes.
You can pass in axis
, keepdim
, and correction
keyword arguments to control the axis along
which the variance is computed, whether the reduced dimensions are retained, and the Bessel's correction applied.
Tensor.manual_seed(42)
t = Tensor.normal(2, 3, mean=2.5, std=0.5)
print(t.numpy())
[[2.9889 2.7339 2.7763]
[2.3356 2.0722 2.6376]]
print(t.var().numpy())
0.109925404
print(t.var(axis=0).numpy())
[0.2134 0.2189 0.0096]
print(t.var(axis=1).numpy())
[0.0187 0.08 ]
Source code in tinygrad/tensor.py
1615 1616 1617 1618 1619 1620 1621 1622 1623 1624 1625 1626 1627 1628 1629 1630 1631 1632 1633 1634 1635 1636 1637 1638 1639 |
|
std
¤
Returns the standard deviation of the tensor along the specified axis or axes.
You can pass in axis
, keepdim
, and correction
keyword arguments to control the axis along
which the standard deviation is computed, whether the reduced dimensions are retained, and the Bessel's correction applied.
Tensor.manual_seed(42)
t = Tensor.normal(2, 3, mean=2.5, std=0.5)
print(t.numpy())
[[2.9889 2.7339 2.7763]
[2.3356 2.0722 2.6376]]
print(t.std().numpy())
0.33155
print(t.std(axis=0).numpy())
[0.462 0.4679 0.0981]
print(t.std(axis=1).numpy())
[0.1367 0.2829]
Source code in tinygrad/tensor.py
1641 1642 1643 1644 1645 1646 1647 1648 1649 1650 1651 1652 1653 1654 1655 1656 1657 1658 1659 1660 1661 1662 1663 |
|
std_mean
¤
Calculates the standard deviation and mean over the dimensions specified by dim.
Syntactic sugar around Tensor.std
and Tensor.mean
to match torch.std_mean
.
Tensor.manual_seed(42)
t = Tensor.normal(2, 3, mean=2.5, std=0.5)
print(t.numpy())
[[2.9889 2.7339 2.7763]
[2.3356 2.0722 2.6376]]
std, mean = t.std_mean()
print(std.numpy(), mean.numpy())
0.33155 2.5907671
Source code in tinygrad/tensor.py
1665 1666 1667 1668 1669 1670 1671 1672 1673 1674 1675 1676 1677 1678 1679 1680 |
|
softmax
¤
softmax(axis=-1, dtype: Optional[DTypeLike] = None)
Applies the softmax function to the tensor along the specified axis.
Rescales the elements of the tensor such that they lie in the range [0, 1] and sum to 1.
You can pass in the axis
keyword argument to control the axis along which the softmax is computed.
Tensor.manual_seed(42)
t = Tensor.randn(2, 3)
print(t.numpy())
[[ 0.9779 0.4678 0.5526]
[-0.3288 -0.8555 0.2753]]
print(t.softmax().numpy())
[[0.4436 0.2664 0.29 ]
[0.2924 0.1727 0.5349]]
print(t.softmax(axis=0).numpy())
[[0.787 0.7897 0.5689]
[0.213 0.2103 0.4311]]
Source code in tinygrad/tensor.py
1688 1689 1690 1691 1692 1693 1694 1695 1696 1697 1698 1699 1700 1701 1702 1703 1704 1705 1706 1707 1708 1709 |
|
log_softmax
¤
log_softmax(axis=-1, dtype: Optional[DTypeLike] = None)
Applies the log-softmax function to the tensor along the specified axis.
The log-softmax function is a numerically stable alternative to the softmax function in log space.
You can pass in the axis
keyword argument to control the axis along which the log-softmax is computed.
Tensor.manual_seed(42)
t = Tensor.randn(2, 3)
print(t.numpy())
[[ 0.9779 0.4678 0.5526]
[-0.3288 -0.8555 0.2753]]
print(t.log_softmax().numpy())
[[-0.8127 -1.3228 -1.238 ]
[-1.2297 -1.7564 -0.6256]]
print(t.log_softmax(axis=0).numpy())
[[-0.2396 -0.2361 -0.564 ]
[-1.5463 -1.5594 -0.8414]]
Source code in tinygrad/tensor.py
1711 1712 1713 1714 1715 1716 1717 1718 1719 1720 1721 1722 1723 1724 1725 1726 1727 1728 1729 1730 1731 1732 |
|
logsumexp
¤
logsumexp(axis=None, keepdim=False)
Computes the log-sum-exp of the tensor along the specified axis or axes.
The log-sum-exp function is a numerically stable way to compute the logarithm of the sum of exponentials.
You can pass in axis
and keepdim
keyword arguments to control the axis along
which the log-sum-exp is computed and whether the reduced dimensions are retained.
Tensor.manual_seed(42)
t = Tensor.randn(2, 3)
print(t.numpy())
[[ 0.9779 0.4678 0.5526]
[-0.3288 -0.8555 0.2753]]
print(t.logsumexp().numpy())
2.1347282
print(t.logsumexp(axis=0).numpy())
[1.2174 0.7039 1.1167]
print(t.logsumexp(axis=1).numpy())
[1.7906 0.9009]
Source code in tinygrad/tensor.py
1734 1735 1736 1737 1738 1739 1740 1741 1742 1743 1744 1745 1746 1747 1748 1749 1750 1751 1752 1753 1754 1755 1756 1757 1758 1759 |
|
logcumsumexp
¤
logcumsumexp(axis=0)
Computes the log-cumsum-exp of the tensor along the specified axis or axes.
The log-cumsum-exp function is a numerically stable way to compute the logarithm of the cumulative sum of exponentials.
You can pass in the axis
keyword argument to control the axis along which
the log-cum-sum-exp is computed.
Tensor.manual_seed(42)
t = Tensor.randn(2, 3)
print(t.numpy())
[[ 0.9779 0.4678 0.5526]
[-0.3288 -0.8555 0.2753]]
print(t.logcumsumexp().numpy())
[[0.9779 0.4678 0.5526]
[1.2174 0.7039 1.1167]]
print(t.logcumsumexp(axis=0).numpy())
[[0.9779 0.4678 0.5526]
[1.2174 0.7039 1.1167]]
print(t.logcumsumexp(axis=1).numpy())
[[ 0.9779 1.4481 1.7906]
[-0.3288 0.1353 0.9009]]
Source code in tinygrad/tensor.py
1761 1762 1763 1764 1765 1766 1767 1768 1769 1770 1771 1772 1773 1774 1775 1776 1777 1778 1779 1780 1781 1782 1783 1784 1785 1786 |
|
argmax
¤
argmax(axis=None, keepdim=False)
Returns the indices of the maximum value of the tensor along the specified axis.
You can pass in axis
and keepdim
keyword arguments to control the axis along
which the maximum is computed and whether the reduced dimensions are retained.
t = Tensor([[1, 0, 2], [5, 4, 3]])
print(t.numpy())
[[1 0 2]
[5 4 3]]
print(t.argmax().numpy()) # Returns the index of the maximum value in the flattened tensor.
3
print(t.argmax(axis=0).numpy()) # Returns the indices of the maximum values along axis 0.
[1 1 1]
print(t.argmax(axis=1).numpy()) # Returns the indices of the maximum values along axis 1.
[2 0]
Source code in tinygrad/tensor.py
1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799 1800 1801 1802 1803 1804 1805 1806 1807 1808 1809 1810 1811 1812 1813 |
|
argmin
¤
argmin(axis=None, keepdim=False)
Returns the indices of the minimum value of the tensor along the specified axis.
You can pass in axis
and keepdim
keyword arguments to control the axis along
which the minimum is computed and whether the reduced dimensions are retained.
t = Tensor([[1, 0, 2], [5, 4, 3]])
print(t.numpy())
[[1 0 2]
[5 4 3]]
print(t.argmin().numpy()) # Returns the index of the minimum value in the flattened tensor.
1
print(t.argmin(axis=0).numpy()) # Returns the indices of the minimum values along axis 0.
[0 0 0]
print(t.argmin(axis=1).numpy()) # Returns the indices of the minimum values along axis 1.
[1 2]
Source code in tinygrad/tensor.py
1815 1816 1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829 1830 1831 1832 1833 1834 1835 1836 |
|
Processing¤
avg_pool2d
¤
avg_pool2d(
kernel_size=(2, 2),
stride=None,
dilation=1,
padding=0,
count_include_pad=True,
)
Applies average pooling over a tensor.
Note
unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.
See: https://paperswithcode.com/method/average-pooling
t = Tensor.arange(25).reshape(1, 1, 5, 5)
print(t.avg_pool2d().numpy())
[[[[ 3. 5.]
[13. 15.]]]]
print(t.avg_pool2d(padding=1).numpy())
[[[[ 0. 0.75 1.75]
[ 3.75 9. 11. ]
[ 8.75 19. 21. ]]]]
Source code in tinygrad/tensor.py
1954 1955 1956 1957 1958 1959 1960 1961 1962 1963 1964 1965 1966 1967 1968 1969 1970 1971 1972 |
|
max_pool2d
¤
max_pool2d(
kernel_size=(2, 2), stride=None, dilation=1, padding=0
)
Applies max pooling over a tensor.
Note
unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.
See: https://paperswithcode.com/method/max-pooling
t = Tensor.arange(25).reshape(1, 1, 5, 5)
print(t.max_pool2d().numpy())
[[[[ 6 8]
[16 18]]]]
print(t.max_pool2d(padding=1).numpy())
[[[[ 0. 2. 4.]
[10. 12. 14.]
[20. 22. 24.]]]]
Source code in tinygrad/tensor.py
1974 1975 1976 1977 1978 1979 1980 1981 1982 1983 1984 1985 1986 1987 1988 1989 1990 1991 |
|
conv2d
¤
conv2d(
weight: Tensor,
bias: Tensor | None = None,
groups=1,
stride=1,
dilation=1,
padding=0,
acc_dtype: DTypeLike | None = None,
) -> Tensor
Applies a convolution over a tensor with a given weight
and optional bias
.
Note
unlike PyTorch, this implementation is not limited to only 2d convolutions and instead works for any number of dimensions.
See: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
t = Tensor.arange(9).reshape(1, 1, 3, 3)
w = Tensor.ones(1, 1, 2, 2)
print(t.conv2d(w).numpy())
[[[[ 8. 12.]
[20. 24.]]]]
Source code in tinygrad/tensor.py
1993 1994 1995 1996 1997 1998 1999 2000 2001 2002 2003 2004 2005 2006 2007 2008 2009 2010 2011 2012 2013 2014 2015 2016 2017 2018 2019 2020 2021 2022 2023 2024 2025 2026 2027 2028 2029 2030 2031 2032 2033 2034 2035 2036 2037 2038 2039 2040 2041 2042 2043 2044 2045 2046 2047 2048 2049 2050 2051 2052 |
|
conv_transpose2d
¤
conv_transpose2d(
weight: Tensor,
bias: Optional[Tensor] = None,
groups=1,
stride=1,
dilation=1,
padding=0,
output_padding=0,
) -> Tensor
Applies a transposed convolution over a tensor with a given weight
and optional bias
.
Note
unlike PyTorch, this implementation is not limited to only 2d transposed convolutions and instead works for any number of dimensions.
See: https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html
t = Tensor.arange(9).reshape(1, 1, 3, 3)
w = Tensor.ones(1, 1, 2, 2)
print(t.conv_transpose2d(w).numpy())
[[[[ 0. 1. 3. 2.]
[ 3. 8. 12. 7.]
[ 9. 20. 24. 13.]
[ 6. 13. 15. 8.]]]]
Source code in tinygrad/tensor.py
2054 2055 2056 2057 2058 2059 2060 2061 2062 2063 2064 2065 2066 2067 2068 2069 2070 2071 2072 2073 2074 2075 2076 2077 2078 |
|
dot
¤
Performs dot product between two tensors.
You can pass in the optional acc_dtype
keyword argument to control the data type of the accumulation.
a = Tensor([[1, 2], [3, 4]])
b = Tensor([[5, 6], [7, 8]])
print(a.dot(b).numpy())
[[19 22]
[43 50]]
Source code in tinygrad/tensor.py
2080 2081 2082 2083 2084 2085 2086 2087 2088 2089 2090 2091 2092 2093 2094 2095 2096 2097 |
|
matmul
¤
Performs matrix multiplication between two tensors.
You can pass in the reverse
keyword argument to control the order of the matrix multiplication.
You can pass in the optional acc_dtype
keyword argument to control the data type of the accumulation.
a = Tensor([[1, 2], [3, 4]])
b = Tensor([[5, 6], [7, 8]])
print(a.matmul(b).numpy())
[[19 22]
[43 50]]
Source code in tinygrad/tensor.py
2099 2100 2101 2102 2103 2104 2105 2106 2107 2108 2109 2110 2111 2112 |
|
einsum
staticmethod
¤
Sums the product of the elements of the input tensors according to a formula based on the Einstein summation convention.
See: https://pytorch.org/docs/stable/generated/torch.einsum.html
x = Tensor([[1, 2], [3, 4]])
y = Tensor([[5, 6], [7, 8]])
print(Tensor.einsum("ij,ij->", x, y).numpy())
70
Source code in tinygrad/tensor.py
1878 1879 1880 1881 1882 1883 1884 1885 1886 1887 1888 1889 1890 1891 1892 1893 1894 1895 1896 1897 1898 1899 1900 1901 1902 1903 1904 1905 1906 1907 1908 1909 1910 1911 1912 1913 1914 1915 1916 1917 1918 1919 1920 1921 1922 |
|
cumsum
¤
Computes the cumulative sum of the tensor along the specified axis.
You can pass in the axis
keyword argument to control the axis along which the cumulative sum is computed.
t = Tensor.ones(2, 3)
print(t.numpy())
[[1. 1. 1.]
[1. 1. 1.]]
print(t.cumsum(1).numpy())
[[1. 2. 3.]
[1. 2. 3.]]
Source code in tinygrad/tensor.py
2118 2119 2120 2121 2122 2123 2124 2125 2126 2127 2128 2129 2130 2131 2132 2133 2134 2135 2136 2137 2138 2139 2140 2141 2142 |
|
triu
¤
Returns the upper triangular part of the tensor, the other elements are set to 0.
The argument diagonal
determines which diagonal is on the boundary. diagonal = 0
means the main diagonal.
Positive diagonal
means above the main diagonal, and negative diagonal
means below the main diagonal.
t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
print(t.numpy())
[[ 1 2 3 4]
[ 5 6 7 8]
[ 9 10 11 12]]
print(t.triu(diagonal=0).numpy())
[[ 1 2 3 4]
[ 0 6 7 8]
[ 0 0 11 12]]
print(t.triu(diagonal=1).numpy())
[[ 0 2 3 4]
[ 0 0 7 8]
[ 0 0 0 12]]
print(t.triu(diagonal=-1).numpy())
[[ 1 2 3 4]
[ 5 6 7 8]
[ 0 10 11 12]]
Source code in tinygrad/tensor.py
2154 2155 2156 2157 2158 2159 2160 2161 2162 2163 2164 2165 2166 2167 2168 2169 2170 2171 2172 2173 2174 2175 |
|
tril
¤
Returns the lower triangular part of the tensor, the other elements are set to 0.
The argument diagonal
determines which diagonal is on the boundary. diagonal = 0
means the main diagonal.
Positive diagonal
means above the main diagonal, and negative diagonal
means below the main diagonal.
t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
print(t.numpy())
[[ 1 2 3 4]
[ 5 6 7 8]
[ 9 10 11 12]]
print(t.tril(diagonal=0).numpy())
[[ 1 0 0 0]
[ 5 6 0 0]
[ 9 10 11 0]]
print(t.tril(diagonal=1).numpy())
[[ 1 2 0 0]
[ 5 6 7 0]
[ 9 10 11 12]]
print(t.tril(diagonal=-1).numpy())
[[ 0 0 0 0]
[ 5 0 0 0]
[ 9 10 0 0]]
Source code in tinygrad/tensor.py
2177 2178 2179 2180 2181 2182 2183 2184 2185 2186 2187 2188 2189 2190 2191 2192 2193 2194 2195 2196 2197 2198 |
|
interpolate
¤
Downsamples or Upsamples to the input size
, accepts 0 to N batch dimensions.
The interpolation algorithm is selected with mode
which currently only supports linear
, nearest
and nearest-exact
.
To run bilinear
or trilinear
, pass in a 2D or 3D size.
t = Tensor([[1, 2, 3, 4], [21, 22, 23, 24], [41, 42, 43, 44]])
print(t.numpy())
[[ 1 2 3 4]
[21 22 23 24]
[41 42 43 44]]
print(t.interpolate(size=(2,3), mode="linear").numpy())
[[ 6 7 8]
[36 37 38]]
Source code in tinygrad/tensor.py
2200 2201 2202 2203 2204 2205 2206 2207 2208 2209 2210 2211 2212 2213 2214 2215 2216 2217 2218 2219 2220 2221 2222 2223 2224 2225 2226 2227 2228 2229 2230 |
|
Neural Network (functional)¤
linear
¤
Applies a linear transformation to self
using weight
and bias
.
See: https://pytorch.org/docs/stable/generated/torch.nn.Linear.html
t = Tensor([[1, 2], [3, 4]])
weight = Tensor([[1, 2], [3, 4]])
bias = Tensor([1, 2])
print(t.linear(weight, bias).numpy())
[[ 8 12]
[16 24]]
Source code in tinygrad/tensor.py
3052 3053 3054 3055 3056 3057 3058 3059 3060 3061 3062 3063 3064 3065 3066 |
|
sequential
¤
Applies a sequence of functions to self
chaining the output of each function to the input of the next.
t = Tensor([1, 2, 3])
print(t.sequential([lambda x: x * 2, lambda x: x + 1]).numpy())
[3 5 7]
Source code in tinygrad/tensor.py
3068 3069 3070 3071 3072 3073 3074 3075 3076 3077 |
|
layernorm
¤
Applies Layer Normalization over a mini-batch of inputs.
- Described: https://paperswithcode.com/method/layer-normalization
- Paper: https://arxiv.org/abs/1607.06450v1
t = Tensor.randn(8, 10, 16) * 2 + 8
print(t.mean().item(), t.std().item())
7.923046112060547 2.0072739124298096
t = t.layernorm()
print(t.mean().item(), t.std().item())
-7.213620367707563e-09 1.0003893375396729
Source code in tinygrad/tensor.py
3079 3080 3081 3082 3083 3084 3085 3086 3087 3088 3089 3090 3091 3092 3093 3094 3095 3096 |
|
batchnorm
¤
batchnorm(
weight: Optional[Tensor],
bias: Optional[Tensor],
mean: Tensor,
invstd: Tensor,
axis: Union[int, Tuple[int, ...]] = 1,
) -> Tensor
Applies Batch Normalization over a mini-batch of inputs.
- Described: https://paperswithcode.com/method/batch-normalization
- Paper: https://arxiv.org/abs/1502.03167
t = Tensor.randn(8, 4, 16, 16) * 2 + 8
print(t.mean().item(), t.std().item())
8.030410766601562 1.9699476957321167
t = t.batchnorm(None, None, t.mean(axis=(0,2,3)), t.var(axis=(0,2,3)).add(1e-5).rsqrt())
print(t.mean().item(), t.std().item())
5.934350610914407e-07 0.9998165369033813
Source code in tinygrad/tensor.py
3098 3099 3100 3101 3102 3103 3104 3105 3106 3107 3108 3109 3110 3111 3112 3113 3114 3115 3116 3117 3118 3119 |
|
dropout
¤
dropout(p=0.5) -> Tensor
Applies dropout to self
.
Note
dropout is only applied when Tensor.training
is True
.
- Described: https://paperswithcode.com/method/dropout
- Paper: https://jmlr.org/papers/v15/srivastava14a.html
Tensor.manual_seed(42)
t = Tensor.randn(2, 2)
with Tensor.train():
print(t.dropout().numpy())
[[ 0. 2.17 ]
[ 0. -0.1682]]
Source code in tinygrad/tensor.py
3121 3122 3123 3124 3125 3126 3127 3128 3129 3130 3131 3132 3133 3134 3135 3136 3137 3138 |
|
one_hot
¤
Converts self
to a one-hot tensor.
num_classes
defaults to -1, which means num_classes will be inferred as max(self) + 1.
t = Tensor([0, 1, 3, 3, 4])
print(t.one_hot(5).numpy())
[[1 0 0 0 0]
[0 1 0 0 0]
[0 0 0 1 0]
[0 0 0 1 0]
[0 0 0 0 1]]
Source code in tinygrad/tensor.py
3140 3141 3142 3143 3144 3145 3146 3147 3148 3149 3150 3151 3152 |
|
scaled_dot_product_attention
¤
scaled_dot_product_attention(
key: Tensor,
value: Tensor,
attn_mask: Optional[Tensor] = None,
dropout_p: float = 0.0,
is_causal: bool = False,
) -> Tensor
Computes scaled dot-product attention.
self
is the query tensor, key
is the key tensor, and value
is the value tensor.
- Described: https://paperswithcode.com/method/scaled
- Paper: https://arxiv.org/abs/1706.03762v7
q = Tensor.randn(2, 4, 8)
k = Tensor.randn(2, 4, 8)
v = Tensor.randn(2, 4, 8)
print(q.scaled_dot_product_attention(k, v).numpy())
[[[-0.1425 -0.1433 -0.3625 0.8853 -0.3129 1.0271 -0.0019 0.2445]
[-0.7137 0.2617 1.1393 0.692 0.0461 0.1132 0.391 -0.3563]
[ 0.4718 0.6791 0.8956 0.9387 -0.7198 0.753 0.5702 0.2661]
[-1.0183 0.005 0.9208 0.6447 0.2658 0.0411 0.2314 -0.4636]]
[[ 0.2928 -0.3364 -0.1937 -0.0755 -0.6196 -0.7339 0.8431 -0.3794]
[ 0.5915 0.3565 -0.6987 0.241 0.2624 -0.1074 -0.3026 -0.3574]
[ 0.3176 -0.4436 -0.3136 -0.5334 -0.5756 -0.851 0.9595 -0.4201]
[ 0.4378 0.0234 -0.0984 0.4847 -0.3579 -0.3998 0.3781 -0.2338]]]
Source code in tinygrad/tensor.py
3154 3155 3156 3157 3158 3159 3160 3161 3162 3163 3164 3165 3166 3167 3168 3169 3170 3171 3172 3173 3174 3175 |
|
binary_crossentropy
¤
Computes the binary cross-entropy loss between self
and Y
.
See: https://pytorch.org/docs/stable/generated/torch.nn.BCELoss.html
t = Tensor([0.1, 0.9, 0.2])
Y = Tensor([0, 1, 0])
print(t.binary_crossentropy(Y).item())
0.14462155103683472
Source code in tinygrad/tensor.py
3182 3183 3184 3185 3186 3187 3188 3189 3190 3191 3192 3193 3194 |
|
binary_crossentropy_logits
¤
Computes the binary cross-entropy loss between self
and Y
where self
is logits.
See: https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html
t = Tensor([-1, 2, -3])
Y = Tensor([0, 1, 0])
print(t.binary_crossentropy_logits(Y).item())
0.16292566061019897
Source code in tinygrad/tensor.py
3196 3197 3198 3199 3200 3201 3202 3203 3204 3205 3206 3207 3208 |
|
sparse_categorical_crossentropy
¤
sparse_categorical_crossentropy(
Y: Tensor,
ignore_index: int = -1,
label_smoothing=0.0,
reduction: ReductionStr = "mean",
) -> Tensor
Computes the sparse categorical cross-entropy loss between self
and Y
.
Note
self
is logits and Y
is the target labels.
NOTE: unlike PyTorch, this function expects the class axis to be -1
See: https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html
t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.sparse_categorical_crossentropy(Y).item())
0.09391524642705917
Source code in tinygrad/tensor.py
3210 3211 3212 3213 3214 3215 3216 3217 3218 3219 3220 3221 3222 3223 3224 3225 3226 3227 3228 3229 3230 3231 3232 3233 |
|
cross_entropy
¤
cross_entropy(
Y: Tensor,
reduction: ReductionStr = "mean",
label_smoothing: float = 0.0,
) -> Tensor
Compute the cross entropy loss between input logits and target.
Note
self
are logits and Y
are the target labels or class probabilities.
See: https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html
t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.cross_entropy(Y).item())
0.09391524642705917
t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.cross_entropy(Y, reduction='none').numpy())
[0.055 0.1328]
Source code in tinygrad/tensor.py
3235 3236 3237 3238 3239 3240 3241 3242 3243 3244 3245 3246 3247 3248 3249 3250 3251 3252 3253 3254 3255 3256 3257 3258 |
|