Creation

Creation (basic)

empty staticmethod

empty(*shape, **kwargs)

Creates an empty tensor with the given shape.

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

t = Tensor.empty(2, 3)
print(t.shape)

(2, 3)

Source code in tinygrad/tensor.py

@staticmethod
def empty(*shape, **kwargs):
  """
  Creates an empty tensor with the given shape.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.empty(2, 3)
  print(t.shape)
  ```
  """
  return Tensor._metaop(MetaOps.EMPTY, argfix(*shape), **kwargs)

zeros staticmethod

zeros(*shape, **kwargs) -> Tensor

Creates a tensor with the given shape, filled with zeros.

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

print(Tensor.zeros(2, 3).numpy())

[[0. 0. 0.]
 [0. 0. 0.]]

print(Tensor.zeros(2, 3, dtype=dtypes.int32).numpy())

[[0 0 0]
 [0 0 0]]

Source code in tinygrad/tensor.py

@staticmethod
def zeros(*shape, **kwargs) -> Tensor:
  """
  Creates a tensor with the given shape, filled with zeros.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.zeros(2, 3).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.zeros(2, 3, dtype=dtypes.int32).numpy())
  ```
  """
  return Tensor.full(argfix(*shape), 0.0, **kwargs)

ones staticmethod

ones(*shape, **kwargs) -> Tensor

Creates a tensor with the given shape, filled with ones.

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

print(Tensor.ones(2, 3).numpy())

[[1. 1. 1.]
 [1. 1. 1.]]

print(Tensor.ones(2, 3, dtype=dtypes.int32).numpy())

[[1 1 1]
 [1 1 1]]

Source code in tinygrad/tensor.py

@staticmethod
def ones(*shape, **kwargs) -> Tensor:
  """
  Creates a tensor with the given shape, filled with ones.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.ones(2, 3).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.ones(2, 3, dtype=dtypes.int32).numpy())
  ```
  """
  return Tensor.full(argfix(*shape), 1.0, **kwargs)

full staticmethod

full(
    shape: Tuple[sint, ...], fill_value: ConstType, **kwargs
) -> Tensor

Creates a tensor with the given shape, filled with the given value.

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

print(Tensor.full((2, 3), 42).numpy())

[[42 42 42]
 [42 42 42]]

print(Tensor.full((2, 3), False).numpy())

[[False False False]
 [False False False]]

Source code in tinygrad/tensor.py

@staticmethod
def full(shape:Tuple[sint, ...], fill_value:ConstType, **kwargs) -> Tensor:
  """
  Creates a tensor with the given shape, filled with the given value.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.full((2, 3), 42).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.full((2, 3), False).numpy())
  ```
  """
  return Tensor(fill_value, **kwargs).reshape((1, )*len(new_shape := argfix(shape))).expand(new_shape)

arange staticmethod

arange(start, stop=None, step=1, **kwargs) -> Tensor

Returns a 1-D tensor of size ceil((stop - start) / step) with values from [start, stop), with spacing between values given by step.

If stop is not specified, values are generated from [0, start) with the given step.

If stop is specified, values are generated from [start, stop) with the given step.

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

print(Tensor.arange(5).numpy())

[0 1 2 3 4]

print(Tensor.arange(5, 10).numpy())

[5 6 7 8 9]

print(Tensor.arange(5, 10, 2).numpy())

[5 7 9]

print(Tensor.arange(5.5, 10, 2).numpy())

[5.5 7.5 9.5]

Source code in tinygrad/tensor.py

@staticmethod
def arange(start, stop=None, step=1, **kwargs) -> Tensor:
  """
  Returns a 1-D tensor of size `ceil((stop - start) / step)` with values from `[start, stop)`, with spacing between values given by `step`.

  If `stop` is not specified, values are generated from `[0, start)` with the given `step`.

  If `stop` is specified, values are generated from `[start, stop)` with the given `step`.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.arange(5).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.arange(5, 10).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.arange(5, 10, 2).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.arange(5.5, 10, 2).numpy())
  ```
  """
  if stop is None: stop, start = start, 0
  assert all(isinstance(s, (int, float)) for s in (start, stop, step)), f"symbolic arange not supported {start=}, {stop=}, {step=}"
  dtype = kwargs.pop("dtype", dtypes.default_float if any(isinstance(x, float) for x in (start, stop, step)) else dtypes.default_int)
  # NOTE: this matches numpy, torch raises RuntimeError if stop-start and step have different signs
  if (stop-start)/step <= 0: return Tensor([], dtype=dtype, **kwargs)
  return (Tensor.full((math.ceil((stop-start)/step),), step, dtype=dtype, **kwargs)._cumsum() + (start - step)).cast(dtype)

eye staticmethod

eye(n: int, m: Optional[int] = None, **kwargs) -> Tensor

Returns a 2-D tensor with n rows and m columns, with ones on the diagonal and zeros elsewhere.

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

print(Tensor.eye(3).numpy())

[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]

print(Tensor.eye(2, 4).numpy())

[[1. 0. 0. 0.]
 [0. 1. 0. 0.]]

Source code in tinygrad/tensor.py

@staticmethod
def eye(n:int, m:Optional[int]=None, **kwargs) -> Tensor:
  """
  Returns a 2-D tensor with `n` rows and `m` columns, with ones on the diagonal and zeros elsewhere.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.eye(3).numpy())
  ```

  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.eye(2, 4).numpy())
  ```
  """
  if n < 0 or (m is not None and m < 0): raise ValueError(f"cannot have negative {n=}, {m=}")
  x = Tensor.ones((n,1),**kwargs).pad((None,(0,n))).flatten().shrink(((0,n*n),)).reshape(n,n)
  return x if m is None else x.pad((None, (0, m-n))) if m > n else x.shrink((None, (0, m)))

full_like

full_like(fill_value: ConstType, **kwargs) -> Tensor

Creates a tensor with the same shape as self, filled with the given value. If dtype is not specified, the dtype of self is used.

You can pass in the device keyword argument to control device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

t = Tensor.ones(2, 3)
print(Tensor.full_like(t, 42).numpy())

[[42. 42. 42.]
 [42. 42. 42.]]

Source code in tinygrad/tensor.py

def full_like(self, fill_value:ConstType, **kwargs) -> Tensor:
  """
  Creates a tensor with the same shape as `self`, filled with the given value.
  If `dtype` is not specified, the dtype of `self` is used.

  You can pass in the `device` keyword argument to control device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.ones(2, 3)
  print(Tensor.full_like(t, 42).numpy())
  ```
  """
  return Tensor.full(self.shape, fill_value, dtype=kwargs.pop("dtype", self.dtype), device=kwargs.pop("device", self.device), **kwargs)

zeros_like

zeros_like(**kwargs) -> Tensor

Creates a tensor with the same shape as self, filled with zeros.

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

t = Tensor.ones(2, 3)
print(Tensor.zeros_like(t).numpy())

[[0. 0. 0.]
 [0. 0. 0.]]

Source code in tinygrad/tensor.py

def zeros_like(self, **kwargs) -> Tensor:
  """
  Creates a tensor with the same shape as `self`, filled with zeros.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.ones(2, 3)
  print(Tensor.zeros_like(t).numpy())
  ```
  """
  return self.full_like(0, **kwargs)

ones_like

ones_like(**kwargs) -> Tensor

Creates a tensor with the same shape as self, filled with ones.

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

t = Tensor.zeros(2, 3)
print(Tensor.ones_like(t).numpy())

[[1. 1. 1.]
 [1. 1. 1.]]

Source code in tinygrad/tensor.py

def ones_like(self, **kwargs) -> Tensor:
  """
  Creates a tensor with the same shape as `self`, filled with ones.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.zeros(2, 3)
  print(Tensor.ones_like(t).numpy())
  ```
  """
  return self.full_like(1, **kwargs)

Creation (random)

manual_seed staticmethod

manual_seed(seed=0)

Sets the seed for random operations.

Tensor.manual_seed(42)
print(Tensor.rand(5).numpy())
print(Tensor.rand(5).numpy())

[0.5053 0.6523 0.4013 0.0438 0.5772]
[0.6669 0.1226 0.8191 0.2581 0.6659]

Tensor.manual_seed(42)  # reset to the same seed
print(Tensor.rand(5).numpy())
print(Tensor.rand(5).numpy())

[0.5053 0.6523 0.4013 0.0438 0.5772]
[0.6669 0.1226 0.8191 0.2581 0.6659]

Source code in tinygrad/tensor.py

@staticmethod
def manual_seed(seed=0):
  """
  Sets the seed for random operations.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  print(Tensor.rand(5).numpy())
  print(Tensor.rand(5).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)  # reset to the same seed
  print(Tensor.rand(5).numpy())
  print(Tensor.rand(5).numpy())
  ```
  """
  Tensor._seed, Tensor._rng_counter = seed, None

rand staticmethod

rand(
    *shape,
    device: Optional[Union[Tuple[str, ...], str]] = None,
    dtype: Optional[DTypeLike] = None,
    **kwargs
) -> Tensor

Creates a tensor with the given shape, filled with random values from a uniform distribution over the interval [0, 1).

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

Tensor.manual_seed(42)
t = Tensor.rand(2, 3)
print(t.numpy())

[[0.5053 0.6523 0.4013]
 [0.0438 0.5772 0.02  ]]

Source code in tinygrad/tensor.py

@staticmethod
def rand(*shape, device:Optional[Union[Tuple[str, ...], str]]=None, dtype:Optional[DTypeLike]=None, **kwargs) -> Tensor:
  """
  Creates a tensor with the given shape, filled with random values from a uniform distribution over the interval `[0, 1)`.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.rand(2, 3)
  print(t.numpy())
  ```
  """
  if not dtypes.is_float(dtype := to_dtype(dtype or dtypes.default_float)): raise ValueError(f"rand only supports float dtypes, got {dtype}")
  if not all_int(shape:=argfix(*shape)) or not all(s >= 0 for s in shape): raise ValueError(f"invalid input {shape=}")
  if (had_counter := Tensor._rng_counter is None): Tensor._rng_counter = Tensor([0], dtype=dtypes.uint32, requires_grad=False)

  if not THREEFRY:
    # for bfloat16, numpy rand passes buffer in float
    if to_dtype(dtype or dtypes.default_float) == dtypes.bfloat16:
      return Tensor.rand(*shape, **kwargs, device=device, dtype=dtypes.float).cast(dtypes.bfloat16)
    return Tensor._metaop(MetaOps.CUSTOM, shape, arg=custom_random, device=device, dtype=dtype, **kwargs)

  # threefry
  if (num := math.ceil(((num_ := prod(shape)) * dtype.itemsize) / 4)) == 0: return Tensor.zeros(shape, device=device, dtype=dtype, **kwargs)
  if not had_counter: Tensor._rng_counter.assign(Tensor._rng_counter + num)
  counts1 = (Tensor.arange(math.ceil(num / 2), device=device, dtype=dtypes.uint32, requires_grad=False)+Tensor._rng_counter.to(device))
  counts2 = counts1 + math.ceil(num / 2)

  # threefry random bits
  x = counts2.cast(dtypes.uint64) << 32 | counts1.cast(dtypes.uint64)
  x = F.Threefry.apply(*x._broadcasted(Tensor._seed))
  counts1, counts2 = (x & 0xffffffff).cast(dtypes.uint32), ((x >> 32) & 0xffffffff).cast(dtypes.uint32)
  bits = counts1.cat(counts2)[:num]

  # bitcast to uint with same number of bits
  _, nmant = dtypes.finfo(dtype)
  uint_dtype = {1: dtypes.uint8, 2: dtypes.uint16, 4: dtypes.uint32, 8: dtypes.uint64}[dtype.itemsize]
  bits = bits.bitcast(uint_dtype)
  # only randomize the mantissa bits and set the exponent to 1
  one = Tensor.ones_like(bits, device=bits.device, dtype=dtype).bitcast(uint_dtype)
  bits = bits.rshift((dtype.itemsize * 8) - nmant).bitwise_or(one)

  # bitcast back to the original dtype
  out = bits.bitcast(dtype)[:num_].sub(1).reshape(shape)
  out.requires_grad = kwargs.get("requires_grad")
  return out.contiguous()

randn staticmethod

randn(
    *shape, dtype: Optional[DTypeLike] = None, **kwargs
) -> Tensor

Creates a tensor with the given shape, filled with random values from a normal distribution with mean 0 and standard deviation 1. If dtype is not specified, the default type is used.

You can pass in the device keyword argument to control device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

Tensor.manual_seed(42)
print(Tensor.randn(2, 3).numpy())

[[-0.8042 -1.1013 -0.9095]
 [ 1.2802 -2.2883  0.7078]]

Source code in tinygrad/tensor.py

@staticmethod
def randn(*shape, dtype:Optional[DTypeLike]=None, **kwargs) -> Tensor:
  """
  Creates a tensor with the given shape, filled with random values from a normal distribution with mean `0` and standard deviation `1`.
  If `dtype` is not specified, the default type is used.

  You can pass in the `device` keyword argument to control device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  print(Tensor.randn(2, 3).numpy())
  ```
  """
  # https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
  src = Tensor.rand((2, *argfix(*shape)), **{**kwargs, "dtype": dtypes.float32})
  return src[0].mul(2*math.pi).cos().mul((1 - src[1]).log().mul(-2).sqrt()).cast(dtype or dtypes.default_float)

randint staticmethod

randint(*shape, low=0, high=10, **kwargs) -> Tensor

Creates a tensor with the given shape, filled with random integer values generated uniformly from the interval [low, high). If dtype is not specified, the default type is used.

You can pass in the device keyword argument to control device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

Tensor.manual_seed(42)
print(Tensor.randint(2, 3, low=5, high=10).numpy())

[[7 8 7]
 [5 7 5]]

Source code in tinygrad/tensor.py

@staticmethod
def randint(*shape, low=0, high=10, **kwargs) -> Tensor:
  """
  Creates a tensor with the given shape, filled with random integer values generated uniformly from the interval `[low, high)`.
  If `dtype` is not specified, the default type is used.

  You can pass in the `device` keyword argument to control device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  print(Tensor.randint(2, 3, low=5, high=10).numpy())
  ```
  """
  if not isinstance(low, int) or not isinstance(high, int): raise TypeError(f"{low=} and {high=} must be integers")
  dtype = kwargs.pop("dtype", dtypes.int32)
  if not dtypes.is_int(dtype): raise TypeError(f"{dtype=} must be int")
  return Tensor.uniform(*shape, low=low, high=high, dtype=dtype, **kwargs)

normal staticmethod

normal(*shape, mean=0.0, std=1.0, **kwargs) -> Tensor

Creates a tensor with the given shape, filled with random values from a normal distribution with the given mean and standard deviation std.

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

Tensor.manual_seed(42)
print(Tensor.normal(2, 3, mean=10, std=2).numpy())

[[ 8.3915  7.7974  8.181 ]
 [12.5603  5.4234 11.4156]]

Source code in tinygrad/tensor.py

@staticmethod
def normal(*shape, mean=0.0, std=1.0, **kwargs) -> Tensor:
  """
  Creates a tensor with the given shape, filled with random values from a normal distribution with the given `mean` and standard deviation `std`.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  print(Tensor.normal(2, 3, mean=10, std=2).numpy())
  ```
  """
  return (std * Tensor.randn(*shape, **kwargs)) + mean

uniform staticmethod

uniform(*shape, low=0.0, high=1.0, **kwargs) -> Tensor

Creates a tensor with the given shape, filled with random values from a uniform distribution over the interval [low, high).

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

Tensor.manual_seed(42)
print(Tensor.uniform(2, 3, low=2, high=10).numpy())

[[6.0426 7.2184 5.2105]
 [2.3502 6.6172 2.1602]]

Source code in tinygrad/tensor.py

@staticmethod
def uniform(*shape, low=0.0, high=1.0, **kwargs) -> Tensor:
  """
  Creates a tensor with the given shape, filled with random values from a uniform distribution over the interval `[low, high)`.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  print(Tensor.uniform(2, 3, low=2, high=10).numpy())
  ```
  """
  dtype = kwargs.pop("dtype", dtypes.default_float)
  return ((high-low) * Tensor.rand(*shape, **kwargs)).cast(dtype) + low

scaled_uniform staticmethod

scaled_uniform(*shape, **kwargs) -> Tensor

Creates a tensor with the given shape, filled with random values from a uniform distribution over the interval [-prod(shape)**-0.5, prod(shape)**-0.5).

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

Tensor.manual_seed(42)
print(Tensor.scaled_uniform(2, 3).numpy())

[[ 0.0043  0.1244 -0.0806]
 [-0.3725  0.063  -0.3919]]

Source code in tinygrad/tensor.py

@staticmethod
def scaled_uniform(*shape, **kwargs) -> Tensor:
  """
  Creates a tensor with the given shape, filled with random values from a uniform distribution
  over the interval `[-prod(shape)**-0.5, prod(shape)**-0.5)`.

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  print(Tensor.scaled_uniform(2, 3).numpy())
  ```
  """
  return Tensor.uniform(*shape, low=-1.0, high=1.0, **kwargs).mul(prod(argfix(*shape))**-0.5)

glorot_uniform staticmethod

glorot_uniform(*shape, **kwargs) -> Tensor

https://www.tensorflow.org/api_docs/python/tf/keras/initializers/GlorotUniform

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

Tensor.manual_seed(42)
print(Tensor.glorot_uniform(2, 3).numpy())

[[ 0.0117  0.3337 -0.2162]
 [-0.9995  0.169  -1.0516]]

Source code in tinygrad/tensor.py

@staticmethod
def glorot_uniform(*shape, **kwargs) -> Tensor:
  """
  <https://www.tensorflow.org/api_docs/python/tf/keras/initializers/GlorotUniform>

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  print(Tensor.glorot_uniform(2, 3).numpy())
  ```
  """
  return Tensor.uniform(*shape, low=-1.0, high=1.0, **kwargs).mul((6/(argfix(*shape)[0]+prod(argfix(*shape)[1:])))**0.5)

kaiming_uniform staticmethod

kaiming_uniform(
    *shape, a: float = 0.01, **kwargs
) -> Tensor

https://pytorch.org/docs/stable/_modules/torch/nn/init.html#kaiming_uniform_

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

Tensor.manual_seed(42)
print(Tensor.kaiming_uniform(2, 3).numpy())

[[ 0.0151  0.4307 -0.2791]
 [-1.2903  0.2182 -1.3575]]

Source code in tinygrad/tensor.py

@staticmethod
def kaiming_uniform(*shape, a:float = 0.01, **kwargs) -> Tensor:
  """
  <https://pytorch.org/docs/stable/_modules/torch/nn/init.html#kaiming_uniform_>

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  print(Tensor.kaiming_uniform(2, 3).numpy())
  ```
  """
  bound = math.sqrt(3.0) * math.sqrt(2.0 / (1 + a ** 2)) / math.sqrt(prod(argfix(*shape)[1:]))
  return Tensor.uniform(*shape, low=-bound, high=bound, **kwargs)

kaiming_normal staticmethod

kaiming_normal(*shape, a: float = 0.01, **kwargs) -> Tensor

https://pytorch.org/docs/stable/_modules/torch/nn/init.html#kaiming_normal_

You can pass in dtype and device keyword arguments to control the data type and device of the tensor. Additionally, all other keyword arguments are passed to the constructor of the tensor.

Tensor.manual_seed(42)
print(Tensor.kaiming_normal(2, 3).numpy())

[[-0.6566 -0.8992 -0.7426]
 [ 1.0452 -1.8683  0.5779]]

Source code in tinygrad/tensor.py

@staticmethod
def kaiming_normal(*shape, a:float = 0.01, **kwargs) -> Tensor:
  """
  <https://pytorch.org/docs/stable/_modules/torch/nn/init.html#kaiming_normal_>

  You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
  Additionally, all other keyword arguments are passed to the constructor of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  print(Tensor.kaiming_normal(2, 3).numpy())
  ```
  """
  std = math.sqrt(2.0 / (1 + a ** 2)) / math.sqrt(prod(argfix(*shape)[1:]))
  return Tensor.normal(*shape, mean=0.0, std=std, **kwargs)