Skip to content

nn (Neural Networks)

Neural Network classes¤

BatchNorm ¤

BatchNorm(
    sz: int,
    eps=1e-05,
    affine=True,
    track_running_stats=True,
    momentum=0.1,
)

Applies Batch Normalization over a 2D or 3D input.

See: Tensor.batchnorm

norm = nn.BatchNorm(3)
t = Tensor.rand(2, 3, 4, 4)
print(t.mean().item(), t.std().item())

0.5637954473495483 0.2746226191520691

t = norm(t)
print(t.mean().item(), t.std().item())

0.5637925863265991 0.27462121844291687

Source code in tinygrad/nn/__init__.py 
32
33
34
35
36
37
38
39
def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
  self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum

  self.weight: Tensor|None = Tensor.ones(sz) if affine else None
  self.bias: Tensor|None = Tensor.zeros(sz) if affine else None

  self.num_batches_tracked = Tensor.zeros(dtype='long', requires_grad=False)
  if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, requires_grad=False), Tensor.ones(sz, requires_grad=False)

Conv1d ¤

Conv1d(
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    stride=1,
    padding: int | str = 0,
    dilation=1,
    groups=1,
    bias=True,
) -> Conv2d

Applies a 1D convolution over an input signal composed of several input planes.

See: https://pytorch.org/docs/stable/generated/torch.nn.Conv1d

conv = nn.Conv1d(1, 1, 3)
t = Tensor.rand(1, 1, 4)
print(t.numpy())

[[[0.9976 0.7907 0.259  0.6945]]]

t = conv(t)
print(t.numpy())

[[[0.6321 0.8225]]]

Source code in tinygrad/nn/__init__.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def Conv1d(in_channels:int, out_channels:int, kernel_size:int, stride=1, padding:int|str=0, dilation=1, groups=1, bias=True) -> Conv2d:
  """
  Applies a 1D convolution over an input signal composed of several input planes.

  See: https://pytorch.org/docs/stable/generated/torch.nn.Conv1d

  ```python exec="true" source="above" session="tensor" result="python"
  conv = nn.Conv1d(1, 1, 3)
  t = Tensor.rand(1, 1, 4)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  t = conv(t)
  print(t.numpy())
  ```
  """
  return Conv2d(in_channels, out_channels, (kernel_size,), stride, padding, dilation, groups, bias)

Conv2d ¤

Conv2d(
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, ...],
    stride=1,
    padding: int | tuple[int, ...] | str = 0,
    dilation=1,
    groups=1,
    bias=True,
)

Applies a 2D convolution over an input signal composed of several input planes.

See: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d

conv = nn.Conv2d(1, 1, 3)
t = Tensor.rand(1, 1, 4, 4)
print(t.numpy())

[[[[0.0757 0.8883 0.0217 0.8138]
   [0.5258 0.8617 0.5675 0.8134]
   [0.6212 0.444  0.7248 0.8333]
   [0.5865 0.333  0.2491 0.2921]]]]

t = conv(t)
print(t.numpy())

[[[[-0.4562 -0.6897]
   [-0.4155 -0.5447]]]]

Source code in tinygrad/nn/__init__.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...], stride=1, padding:int|tuple[int, ...]|str=0,
             dilation=1, groups=1, bias=True):
  self.kernel_size = make_tuple(kernel_size, 2)
  if isinstance(padding, str):
    if padding.lower() != 'same': raise ValueError(f"Invalid padding string {padding!r}, only 'same' is supported")
    if stride != 1: raise ValueError("padding='same' is not supported for strided convolutions")
    pad = [(d*(k-1)//2, d*(k-1) - d*(k-1)//2) for d,k in zip(make_tuple(dilation, len(self.kernel_size)), self.kernel_size[::-1])]
    padding = tuple(flatten(pad))
  self.stride, self.dilation, self.groups, self.padding = stride, dilation, groups, padding
  scale = 1 / math.sqrt(in_channels * prod(self.kernel_size))
  self.weight = Tensor.uniform(out_channels, in_channels//groups, *self.kernel_size, low=-scale, high=scale)
  self.bias: Tensor|None = Tensor.uniform(out_channels, low=-scale, high=scale) if bias else None

ConvTranspose1d ¤

ConvTranspose1d(
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    stride=1,
    padding=0,
    output_padding=0,
    dilation=1,
    groups=1,
    bias=True,
) -> ConvTranspose2d

Applies a 1D transposed convolution operator over an input signal composed of several input planes.

See: https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose1d

conv = nn.ConvTranspose1d(1, 1, 3)
t = Tensor.rand(1, 1, 4)
print(t.numpy())

[[[0.4502 0.8436 0.16   0.5108]]]

t = conv(t)
print(t.numpy())

[[[0.6609 0.8448 0.6035 0.7054 0.5158 0.4887]]]

Source code in tinygrad/nn/__init__.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def ConvTranspose1d(in_channels:int, out_channels:int, kernel_size:int, stride=1, padding=0, output_padding=0, dilation=1,
                      groups=1, bias=True) -> ConvTranspose2d:
  """
  Applies a 1D transposed convolution operator over an input signal composed of several input planes.

  See: https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose1d

  ```python exec="true" source="above" session="tensor" result="python"
  conv = nn.ConvTranspose1d(1, 1, 3)
  t = Tensor.rand(1, 1, 4)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  t = conv(t)
  print(t.numpy())
  ```
  """
  return ConvTranspose2d(in_channels, out_channels, (kernel_size,), stride, padding, output_padding, dilation, groups, bias)

ConvTranspose2d ¤

ConvTranspose2d(
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, ...],
    stride=1,
    padding=0,
    output_padding=0,
    dilation=1,
    groups=1,
    bias=True,
)

Bases: Conv2d

Applies a 2D transposed convolution operator over an input image.

See: https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d

conv = nn.ConvTranspose2d(1, 1, 3)
t = Tensor.rand(1, 1, 4, 4)
print(t.numpy())

[[[[0.4952 0.9488 0.6573 0.8941]
   [0.5037 0.0963 0.6511 0.26  ]
   [0.5433 0.0289 0.9932 0.615 ]
   [0.1001 0.8292 0.7317 0.831 ]]]]

t = conv(t)
print(t.numpy())

[[[[-0.0021 -0.1133 -0.1806 -0.1092 -0.1486  0.0428]
   [ 0.0585 -0.0172  0.1394  0.196   0.0944  0.266 ]
   [ 0.0449 -0.2152  0.0477 -0.3644 -0.0702  0.0448]
   [ 0.0594 -0.1311  0.0339 -0.2116  0.0204  0.1932]
   [ 0.0061  0.0367 -0.0037  0.1457  0.008   0.1979]
   [ 0.0062 -0.0262 -0.1181 -0.1575 -0.1399 -0.0462]]]]

Source code in tinygrad/nn/__init__.py 
146
147
148
149
150
151
def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...], stride=1, padding=0, output_padding=0,
              dilation=1, groups=1, bias=True):
  super().__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
  scale = 1 / math.sqrt(in_channels * prod(self.kernel_size))
  self.weight = Tensor.uniform(in_channels, out_channels//groups, *self.kernel_size, low=-scale, high=scale)
  self.output_padding = output_padding

Linear ¤

Linear(in_features: int, out_features: int, bias=True)

Applies a linear transformation to the incoming data.

See: https://pytorch.org/docs/stable/generated/torch.nn.Linear

lin = nn.Linear(3, 4)
t = Tensor.rand(2, 3)
print(t.numpy())

[[0.7369 0.357  0.552 ]
 [0.7649 0.8311 0.6705]]

t = lin(t)
print(t.numpy())

[[-0.1068  0.76    0.755  -0.323 ]
 [-0.3972  0.974   0.6014 -0.2679]]

Source code in tinygrad/nn/__init__.py 
172
173
174
175
def __init__(self, in_features:int, out_features:int, bias=True):
  bound = 1 / math.sqrt(in_features)
  self.weight = Tensor.uniform(out_features, in_features, low=-bound, high=bound)
  self.bias = Tensor.uniform(out_features, low=-bound, high=bound) if bias else None

GroupNorm ¤

GroupNorm(
    num_groups: int,
    num_channels: int,
    eps=1e-05,
    affine=True,
)

Applies Group Normalization over a mini-batch of inputs.

norm = nn.GroupNorm(2, 12)
t = Tensor.rand(2, 12, 4, 4) * 2 + 1
print(t.mean().item(), t.std().item())

2.0552477836608887 0.5583968162536621

t = norm(t)
print(t.mean().item(), t.std().item())

-1.584819386835079e-07 1.0012880563735962

Source code in tinygrad/nn/__init__.py 
195
196
197
198
def __init__(self, num_groups:int, num_channels:int, eps=1e-5, affine=True):
  self.num_groups, self.num_channels, self.eps = num_groups, num_channels, eps
  self.weight: Tensor|None = Tensor.ones(num_channels) if affine else None
  self.bias: Tensor|None = Tensor.zeros(num_channels) if affine else None

InstanceNorm ¤

InstanceNorm(
    num_features: int,
    eps: float = 1e-05,
    affine: bool = True,
)

Applies Instance Normalization over a mini-batch of inputs.

norm = nn.InstanceNorm(3)
t = Tensor.rand(2, 3, 4, 4) * 2 + 1
print(t.mean().item(), t.std().item())

1.9553427696228027 0.5989101529121399

t = norm(t)
print(t.mean().item(), t.std().item())

-2.892409156629583e-08 1.0052343606948853

Source code in tinygrad/nn/__init__.py 
225
226
227
228
def __init__(self, num_features:int, eps:float=1e-5, affine:bool=True):
  self.num_features, self.eps = num_features, eps
  self.weight: Tensor|None = Tensor.ones(num_features) if affine else None
  self.bias: Tensor|None = Tensor.zeros(num_features) if affine else None

LayerNorm ¤

LayerNorm(
    normalized_shape: int | tuple[int, ...],
    eps: float = 1e-05,
    elementwise_affine: bool = True,
)

Applies Layer Normalization over a mini-batch of inputs.

norm = nn.LayerNorm(3)
t = Tensor.rand(2, 5, 3) * 2 + 1
print(t.mean().item(), t.std().item())

2.1495094299316406 0.5331981778144836

t = norm(t)
print(t.mean().item(), t.std().item())

-4.4447091340771294e-07 1.01702880859375

Source code in tinygrad/nn/__init__.py 
251
252
253
254
255
def __init__(self, normalized_shape:int|tuple[int, ...], eps:float=1e-5, elementwise_affine:bool=True):
  self.normalized_shape: tuple[int, ...] = make_tuple(normalized_shape, 1)
  self.axis, self.eps = tuple(-1-i for i in range(len(self.normalized_shape))), eps
  self.weight: Tensor|None = Tensor.ones(*self.normalized_shape) if elementwise_affine else None
  self.bias: Tensor|None = Tensor.zeros(*self.normalized_shape) if elementwise_affine else None

LayerNorm2d ¤

LayerNorm2d(
    normalized_shape: int | tuple[int, ...],
    eps: float = 1e-05,
    elementwise_affine: bool = True,
)

Bases: LayerNorm

Applies Layer Normalization over a mini-batch of 2D inputs.

See: LayerNorm

norm = nn.LayerNorm2d(3)
t = Tensor.rand(2, 3, 4, 4) * 2 + 1
print(t.mean().item(), t.std().item())

1.995133638381958 0.5704997181892395

t = norm(t)
print(t.mean().item(), t.std().item())

-2.039905950823595e-07 1.0051875114440918

Source code in tinygrad/nn/__init__.py 
251
252
253
254
255
def __init__(self, normalized_shape:int|tuple[int, ...], eps:float=1e-5, elementwise_affine:bool=True):
  self.normalized_shape: tuple[int, ...] = make_tuple(normalized_shape, 1)
  self.axis, self.eps = tuple(-1-i for i in range(len(self.normalized_shape))), eps
  self.weight: Tensor|None = Tensor.ones(*self.normalized_shape) if elementwise_affine else None
  self.bias: Tensor|None = Tensor.zeros(*self.normalized_shape) if elementwise_affine else None

RMSNorm ¤

RMSNorm(dim: int, eps=1e-06, elementwise_affine=True)

Applies Root Mean Square Normalization to input.

norm = nn.RMSNorm(4)
t = Tensor.arange(12, dtype=dtypes.float).reshape(3, 4)
print(t.numpy())

[[ 0.  1.  2.  3.]
 [ 4.  5.  6.  7.]
 [ 8.  9. 10. 11.]]

print(norm(t).numpy())

[[0.     0.5345 1.069  1.6036]
 [0.7127 0.8909 1.069  1.2472]
 [0.8363 0.9409 1.0454 1.15  ]]

Source code in tinygrad/nn/__init__.py 
296
297
298
def __init__(self, dim:int, eps=1e-6, elementwise_affine=True):
  self.eps = eps
  self.weight = Tensor.ones(dim) if elementwise_affine else None

Embedding ¤

Embedding(vocab_size: int, embed_size: int)

A simple lookup table that stores embeddings of a fixed dictionary and size.

See: https://pytorch.org/docs/stable/generated/torch.nn.Embedding

emb = nn.Embedding(10, 3)
print(emb(Tensor([1, 2, 3, 1])).numpy())

[[ 0.3989  0.0204  0.2396]
 [ 0.0843  0.4843  0.5914]
 [-0.1359 -0.6479  0.6577]
 [ 0.3989  0.0204  0.2396]]
Source code in tinygrad/nn/__init__.py 
357
358
def __init__(self, vocab_size:int, embed_size:int):
  self.weight = Tensor.glorot_uniform(vocab_size, embed_size)

LSTMCell ¤

LSTMCell(
    input_size: int, hidden_size: int, bias: bool = True
)

A long short-term memory (LSTM) cell.

Parameters:

  • input_size (int) –

    The number of expected features in the input x

  • hidden_size (int) –

    The number of features in the hidden state h

  • bias (bool, default: True ) –

    If False, then the layer does not use bias weights b_ih and b_hh

Source code in tinygrad/nn/__init__.py 
374
375
376
377
378
379
def __init__(self, input_size:int, hidden_size:int, bias:bool=True):
  stdv = 1.0 / math.sqrt(hidden_size)
  self.weight_ih = Tensor.uniform(hidden_size*4, input_size, low=-stdv, high=stdv)
  self.weight_hh = Tensor.uniform(hidden_size*4, hidden_size, low=-stdv, high=stdv)
  self.bias_ih: Tensor|None = Tensor.zeros(hidden_size*4) if bias else None
  self.bias_hh: Tensor|None = Tensor.zeros(hidden_size*4) if bias else None

Optimizers¤

SGD ¤

SGD(
    params: list[Tensor],
    lr=0.001,
    momentum=0.0,
    weight_decay=0.0,
    nesterov=False,
    classic=False,
    device=None,
    fused=FUSE_OPTIM,
)

Stochastic Gradient Descent (SGD) optimizer with optional momentum and weight decay.

classic is a boolean flag that determines whether to use the popular momentum update rule or the classic momentum update rule.

Source code in tinygrad/nn/optim.py 
77
78
79
80
81
82
83
def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov=False, classic=False, device=None, fused=FUSE_OPTIM):
  """
  Stochastic Gradient Descent (SGD) optimizer with optional momentum and weight decay.

  `classic` is a boolean flag that determines whether to use the popular momentum update rule or the classic momentum update rule.
  """
  return LARS(params, lr, momentum, weight_decay, 0, None, nesterov, classic=classic, pre_wd=True, tcoef=0.0, device=device, fused=fused)

LARS ¤

LARS(
    params: list[Tensor],
    lr=0.001,
    momentum=0.9,
    weight_decay=0.0001,
    ns_steps=0,
    ns_coefficients=None,
    nesterov=False,
    classic=True,
    pre_wd=True,
    tcoef=0.001,
    device=None,
    fused=FUSE_OPTIM,
)

Bases: Optimizer

Layer-wise Adaptive Rate Scaling (LARS) optimizer with optional momentum and weight decay.

Source code in tinygrad/nn/optim.py 
104
105
106
107
108
109
def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, ns_steps=0, ns_coefficients=None,
             nesterov=False, classic=True, pre_wd=True, tcoef=0.001, device=None, fused=FUSE_OPTIM):
  super().__init__(params, lr, device, fused)
  self.momentum, self.wd, self.ns_steps, self.ns_coefficients  = momentum, weight_decay, ns_steps, ns_coefficients
  self.nesterov, self.classic, self.pre_wd, self.tcoef = nesterov, classic, pre_wd, tcoef
  self.b = self._new_optim_param() if self.momentum else []

AdamW ¤

AdamW(
    params: list[Tensor],
    lr=0.001,
    b1=0.9,
    b2=0.999,
    eps=1e-08,
    weight_decay=0.01,
    device=None,
    fused=FUSE_OPTIM,
)

AdamW optimizer with optional weight decay.

Source code in tinygrad/nn/optim.py 
134
135
136
137
138
139
140
def AdamW(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.01, device=None, fused=FUSE_OPTIM):
  """
  AdamW optimizer with optional weight decay.

  - Paper: https://arxiv.org/abs/1711.05101v3
  """
  return LAMB(params, lr, b1, b2, eps, weight_decay, adam=True, device=device, fused=fused)

Adam ¤

Adam(
    params: list[Tensor],
    lr=0.001,
    b1=0.9,
    b2=0.999,
    eps=1e-08,
    device=None,
    fused=FUSE_OPTIM,
)

Adam optimizer.

Source code in tinygrad/nn/optim.py 
141
142
143
144
145
146
147
def Adam(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, device=None, fused=FUSE_OPTIM):
  """
  Adam optimizer.

  - Paper: https://arxiv.org/abs/1412.6980
  """
  return LAMB(params, lr, b1, b2, eps, 0.0, adam=True, device=device, fused=fused)

LAMB ¤

LAMB(
    params: list[Tensor],
    lr=0.001,
    b1=0.9,
    b2=0.999,
    eps=1e-06,
    weight_decay=0.0,
    adam=False,
    device=None,
    fused=FUSE_OPTIM,
)

Bases: Optimizer

LAMB optimizer with optional weight decay.

Source code in tinygrad/nn/optim.py 
155
156
157
158
159
160
def __init__(self, params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, adam=False, device=None, fused=FUSE_OPTIM):
  super().__init__(params, lr, device, fused)
  self.b1, self.b2, self.eps, self.wd, self.adam = b1, b2, eps, weight_decay, adam
  self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) for _ in [b1, b2])
  self.m = self._new_optim_param()
  self.v = self._new_optim_param()

Load/Save¤

safe_load ¤

safe_load(fn: Tensor | str | Path) -> dict[str, Tensor]

Loads a .safetensor file, returning the state_dict.

state_dict = nn.state.safe_load("test.safetensor")
Source code in tinygrad/nn/state.py 
50
51
52
53
54
55
56
57
58
59
60
61
def safe_load(fn:Tensor|str|pathlib.Path) -> dict[str, Tensor]:
  """
  Loads a .safetensor file, returning the `state_dict`.

  ```python
  state_dict = nn.state.safe_load("test.safetensor")
  ```
  """
  t, data_start, metadata = safe_load_metadata(fn)
  data = t[data_start:]
  return { k: data[v['data_offsets'][0]:v['data_offsets'][1]].bitcast(safe_dtypes[v['dtype']]).reshape(v['shape'])
          for k, v in metadata.items() if k != "__metadata__" }

safe_save ¤

safe_save(
    tensors: dict[str, Tensor],
    fn: str,
    metadata: dict[str, Any] | None = None,
)

Saves a state_dict to disk in a .safetensor file with optional metadata.

t = Tensor([1, 2, 3])
nn.state.safe_save({'t':t}, "test.safetensor")
Source code in tinygrad/nn/state.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=None):
  """
  Saves a `state_dict` to disk in a .safetensor file with optional metadata.

  ```python
  t = Tensor([1, 2, 3])
  nn.state.safe_save({'t':t}, "test.safetensor")
  ```
  """
  headers, offset = {}, 0
  if metadata: headers['__metadata__'] = metadata
  for k,v in tensors.items():
    headers[k] = {'dtype': inverse_safe_dtypes[v.dtype], 'shape': list(v.shape), 'data_offsets':[offset, offset+v.nbytes()]}
    offset += v.nbytes()
  j = json.dumps(headers, separators=(',', ':'))
  j += "\x20"*(round_up(len(j),8)-len(j))
  pathlib.Path(fn).unlink(missing_ok=True)
  t = Tensor.empty(8+len(j)+offset, dtype=dtypes.uint8, device=f"disk:{fn}")
  t[0:8].bitcast(dtypes.int64).assign([len(j)])
  t[8:8+len(j)].assign(list(j.encode('utf-8')))
  for k,v in safe_load(t).items(): v.assign(tensors[k])

get_state_dict ¤

get_state_dict(
    obj, prefix: str = "", tensor_type=Tensor
) -> dict[str, Tensor]

Returns a state_dict of the object, with optional prefix.

class Net:
  def __init__(self):
    self.l1 = nn.Linear(4, 5)
    self.l2 = nn.Linear(5, 6)

net = Net()
print(nn.state.get_state_dict(net).keys())

dict_keys(['l1.weight', 'l1.bias', 'l2.weight', 'l2.bias'])
Source code in tinygrad/nn/state.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def get_state_dict(obj, prefix:str='', tensor_type=Tensor) -> dict[str, Tensor]:
  """
  Returns a `state_dict` of the object, with optional prefix.

  ```python exec="true" source="above" session="tensor" result="python"
  class Net:
    def __init__(self):
      self.l1 = nn.Linear(4, 5)
      self.l2 = nn.Linear(5, 6)

  net = Net()
  print(nn.state.get_state_dict(net).keys())
  ```
  """
  if isinstance(obj, tensor_type): return {prefix.strip('.'):obj}
  if hasattr(obj, '_asdict'): return get_state_dict(obj._asdict(), prefix, tensor_type)  # namedtuple
  if isinstance(obj, OrderedDict): return get_state_dict(dict(obj), prefix, tensor_type)
  if hasattr(obj, '__dict__'): return get_state_dict(obj.__dict__, prefix, tensor_type)
  state_dict = {}
  if isinstance(obj, (list, tuple)):
    for i,x in enumerate(obj): state_dict.update(get_state_dict(x, f"{prefix}{str(i)}.", tensor_type))
  elif isinstance(obj, dict):
    for k,v in obj.items(): state_dict.update(get_state_dict(v, f"{prefix}{str(k)}.", tensor_type))
  return state_dict

get_parameters ¤

get_parameters(obj) -> list[Tensor]

class Net:
  def __init__(self):
    self.l1 = nn.Linear(4, 5)
    self.l2 = nn.Linear(5, 6)

net = Net()
print(len(nn.state.get_parameters(net)))

4
Source code in tinygrad/nn/state.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
def get_parameters(obj) -> list[Tensor]:
  """
  ```python exec="true" source="above" session="tensor" result="python"
  class Net:
    def __init__(self):
      self.l1 = nn.Linear(4, 5)
      self.l2 = nn.Linear(5, 6)

  net = Net()
  print(len(nn.state.get_parameters(net)))
  ```
  """
  return list(get_state_dict(obj).values())

load_state_dict ¤

load_state_dict(
    model,
    state_dict: dict[str, Tensor],
    strict=True,
    verbose=True,
    consume=False,
    realize=True,
) -> list[Tensor]

Loads a state_dict into a model. Return the loaded Tensors.

class Net:
  def __init__(self):
    self.l1 = nn.Linear(4, 5)
    self.l2 = nn.Linear(5, 6)

net = Net()
state_dict = nn.state.get_state_dict(net)
nn.state.load_state_dict(net, state_dict)
Source code in tinygrad/nn/state.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=True, consume=False, realize=True) -> list[Tensor]:
  """
  Loads a `state_dict` into a model. Return the loaded Tensors.

  ```python
  class Net:
    def __init__(self):
      self.l1 = nn.Linear(4, 5)
      self.l2 = nn.Linear(5, 6)

  net = Net()
  state_dict = nn.state.get_state_dict(net)
  nn.state.load_state_dict(net, state_dict)
  ```
  """
  start_mem_used = GlobalCounters.mem_used
  ret = []
  with Timing("loaded weights in ",
              lambda et_ns: f", {(B:=(GlobalCounters.mem_used-start_mem_used))/1e9:.2f} GB loaded at {B/et_ns:.2f} GB/s", enabled=verbose):
    model_state_dict = get_state_dict(model)
    if DEBUG >= 1 and len(state_dict) > len(model_state_dict):
      print("WARNING: unused weights in state_dict", sorted(list(state_dict.keys() - model_state_dict.keys())))
    for k,v in (t := tqdm(model_state_dict.items(), disable=CI or not verbose)):
      t.desc = f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, {k:50s}: "
      if k not in state_dict and not strict:
        if DEBUG >= 1: print(f"WARNING: not loading {k}")
        continue
      if v.shape != state_dict[k].shape:
        if {(), (1,)} == {state_dict[k].shape, v.shape}: state_dict[k] = state_dict[k].reshape(v.shape)
        else: raise ValueError(f'Shape mismatch in layer `{k}`: Expected shape {v.shape}, but found {state_dict[k].shape} in state dict.')
      if isinstance(v.device, tuple):
        if isinstance(state_dict[k].device, tuple): v.replace(state_dict[k])
        else: v.replace(state_dict[k].shard(v.device, v.uop.axis))
      else: v.replace(state_dict[k].to(v.device))
      if realize: v.realize()
      if consume: del state_dict[k]
      ret.append(v)
  return ret

tar_extract ¤

tar_extract(fn: Tensor | str | Path) -> dict[str, Tensor]

Extracts files from a tar archive and returns them as a dictionary of names (keys) and tensors (values).

tensors = nn.state.tar_extract(Tensor(pathlib.Path("archive.tar")))
Source code in tinygrad/nn/state.py 
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
@accept_filename
def tar_extract(t: Tensor) -> dict[str, Tensor]:
  """
  ```python
  tar_extract(fn: Tensor | str | Path) -> dict[str, Tensor]
  ```

  Extracts files from a tar archive and returns them as a dictionary of names (keys) and tensors (values).

  ```python
  tensors = nn.state.tar_extract(Tensor(pathlib.Path("archive.tar")))
  ```
  """
  with tarfile.open(fileobj=TensorIO(t), mode="r") as tar:
    return {member.name:t[member.offset_data:member.offset_data+member.size] for member in tar if member.type == tarfile.REGTYPE}

torch_load ¤

torch_load(fn: Tensor | str | Path) -> dict[str, Tensor]

Loads a torch .pth file, returning the state_dict.

state_dict = nn.state.torch_load("test.pth")
Source code in tinygrad/nn/state.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
@accept_filename
def torch_load(t:Tensor) -> dict[str, Tensor]:
  """
  ```python
  torch_load(fn: Tensor | str | Path) -> dict[str, Tensor]
  ```

  Loads a torch .pth file, returning the `state_dict`.

  ```python
  state_dict = nn.state.torch_load("test.pth")
  ```
  """
  storage_source: dict[str|int, Tensor] = {}
  lens: dict[str|int, int] = {}

  def _rebuild_tensor(storage, storage_offset, size, stride):
    return _rebuild_tensor_v2(storage, storage_offset, size, stride)

  def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad=None, backward_hooks=None, metadata=None):
    #print(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata)
    lens[storage[2]] = storage[4] * storage[1].itemsize
    if storage[2] not in storage_source: return None
    byte_start, byte_end = storage_offset*storage[1].itemsize, (storage_offset + prod(size))*storage[1].itemsize
    ret = storage_source[storage[2]][byte_start:byte_end].bitcast(storage[1])

    # 7 lines to deal with permuted tensors. NOTE: this currently requires reading off the disk
    shape_strides = [(s, st) for s,st in zip(size, stride) if s != 1]
    permute_indexes = [len(shape_strides)-1-y for y in argsort([x[1] for x in shape_strides])]
    if tuple(permute_indexes) != tuple(range(len(permute_indexes))):
      intermediate_shape = tuple([shape_strides[x][0] for x in argsort(permute_indexes)])
      assert tuple([shape_strides[i][1] for i in argsort(permute_indexes)]) == strides_for_shape(intermediate_shape), "nonpermutable strides"
      if DEBUG >= 3: print(f"WARNING: this torch load is slow. to permute {intermediate_shape} with {permute_indexes}")
      assert storage[1] != dtypes.bfloat16, "can't permute BF16"
      # TODO: find a nice way to support all movement ops on disktensors
      ret = ret.to(None).reshape(intermediate_shape).permute(permute_indexes)

    return ret.reshape(size)

  class Parameter:
    def __setstate__(self, state): self.tensor = state[0]

  deserialized_objects: dict[str, Any] = {}
  intercept = {"HalfStorage": dtypes.float16, "FloatStorage": dtypes.float32, "BFloat16Storage": dtypes.bfloat16,
               "IntStorage": dtypes.int32, "BoolStorage": dtypes.bool,
               "LongStorage": dtypes.int64, "_rebuild_tensor": _rebuild_tensor, "_rebuild_tensor_v2": _rebuild_tensor_v2,
               "FloatTensor": None, "Parameter": Parameter}
  whitelist = {"torch", "collections", "numpy", "_codecs"}  # NOTE: this is not for security, only speed
  class Dummy: pass
  class TorchPickle(pickle.Unpickler):
    def find_class(self, module, name):
      module_root = module.split(".")[0]
      if module_root not in whitelist:
        if DEBUG >= 2: print(f"WARNING: returning Dummy for {module} {name}")
        return Dummy
      return intercept[name] if module_root == "torch" else super().find_class(module, name)
    def persistent_load(self, pid): return deserialized_objects.get(pid, pid)

  fobj = io.BufferedReader(TensorIO(t))
  def passthrough_reset(v: bool): return fobj.seek(0, 0) or v
  if passthrough_reset(zipfile.is_zipfile(fobj)): # NOTE: passthrough_reset required to support python < 3.14
    files = zip_extract(t)
    base_name = next(iter(files)).split('/', 1)[0]
    # keyed by persistent_id in pickle file
    storage_source = {fn.split("/")[-1]: data for fn, data in files.items() if fn.startswith(f"{base_name}/data/") and not fn.endswith(".pkl")}
    return TorchPickle(io.BufferedReader(TensorIO(files[f"{base_name}/data.pkl"]), 1_000_000)).load()
  elif passthrough_reset(tarfile.is_tarfile(fobj)): # NOTE: passthrough_reset required to support python < 3.11
    files = tar_extract(t)
    f = io.BufferedReader(TensorIO(files["storages"]), 1_000_000)
    # slice source tensor t
    for _ in range(TorchPickle(f).load()):
      (key, _, storage_type), sz = TorchPickle(f).load(), struct.unpack('<q', f.read(8))[0]
      byte_offset = f.tell()
      storage_source[key] = files["storages"][byte_offset:byte_offset + sz * storage_type.itemsize]
      f.seek(sz * storage_type.itemsize, 1)
    f = io.BufferedReader(TensorIO(files["tensors"]), 1_000_000)
    # get tensor metadata
    for _ in range(TorchPickle(f).load()):
      (key, storage_id, _), ndim, _ = TorchPickle(f).load(), struct.unpack('<i', f.read(4))[0], f.read(4)
      size, stride = struct.unpack(f'<{ndim}q', f.read(8 * ndim)), struct.unpack(f'<{ndim}q', f.read(8 * ndim))
      storage_offset = struct.unpack('<q', f.read(8))[0]
      deserialized_objects[str(key)] = _rebuild_tensor_v2((None, storage_type, storage_id, None, -1), storage_offset, size, stride)
    pkl_data = TorchPickle(io.BufferedReader(TensorIO(files["pickle"]), 1_000_000)).load()
    return {k: v.tensor if isinstance(v, Parameter) else v for k, v in pkl_data.items()}
  else:
    pkl = TorchPickle(fobj)
    _, _, _, rwd, _, ids, base_offset = pkl.load(), pkl.load(), pkl.load(), fobj.tell(), pkl.load(), pkl.load(), fobj.tell()
    # slice source tensor t
    for i in ids:
      storage_source[i] = t[base_offset + 8:base_offset + 8 + lens[i]]
      base_offset += 8 + lens[i]
    fobj.seek(rwd)
    return TorchPickle(fobj).load()

gguf_load ¤

gguf_load(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]

Loads a .gguf file, returning the kv_data and state_dict.

gguf_tensor = Tensor(pathlib.Path("Meta-Llama-3-8B-Instruct.Q4_0.gguf")).to(Device.DEFAULT)
kv_data, state_dict = nn.state.gguf_load(gguf_tensor)

Note

The provided tensor must be on a device that supports execution.

Source code in tinygrad/nn/state.py 
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
@accept_filename
def gguf_load(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
  """
  Loads a .gguf file, returning the `kv_data` and `state_dict`.

  ```python
  gguf_tensor = Tensor(pathlib.Path("Meta-Llama-3-8B-Instruct.Q4_0.gguf")).to(Device.DEFAULT)
  kv_data, state_dict = nn.state.gguf_load(gguf_tensor)
  ```

  NOTE: The provided tensor must be on a device that supports execution.
  """
  reader, kv_data, state_dict = io.BufferedReader(TensorIO(tensor), 1_000_000), {}, {}
  def read_unpack(fmt: str, n: int): return struct.unpack(fmt, reader.read(n))[0]
  def read_str(): return str(reader.read(read_uint64()), "utf-8")
  def read_arr():
    reader, n = readers[read_int32()], read_uint64()
    return [ reader() for _ in range(n) ]

  readers: dict[int, Callable[[], Any]] = { 8: read_str, 9: read_arr, **{ t: functools.partial(read_unpack, "<"+f, nb) for t,f,nb in \
    [ (0,"c",1), (1,"b",1), (2,"H",2), (3,"h",2), (4,"I",4), (5,"i",4), (6,"f",4), (7,"?",1), (10,"Q",8), (11,"q",8), (12,"d",8) ] } }
  read_uint32, read_int32, read_uint64, read_int64 = readers[4], readers[5], readers[10], readers[11]

  magic, version, n_tensors, n_kv = reader.read(4), read_int32(), read_int64(), read_int64()
  if magic != b"GGUF" or version not in [2, 3]: raise ValueError("Invalid GGUF format!")
  for _ in range(n_kv):
    k, typ = read_str(), read_int32()
    kv_data[k] = readers[typ]()

  t_infos = [ (read_str(), tuple(read_uint64() for _ in range(read_uint32())), read_int32(), read_uint64()) for _ in range(n_tensors) ]
  alignment, pos = kv_data.get("general.alignment", 32), reader.tell()
  data_start = round_up(pos, alignment)

  for name, dims, typ, off in t_infos: state_dict[name] = ggml_data_to_tensor(tensor[data_start + off:], prod(dims), typ).reshape(*reversed(dims))

  return kv_data, state_dict