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.

• Described: https://paperswithcode.com/method/batch-normalization
• Paper: https://arxiv.org/abs/1502.03167v3

See: Tensor.batchnorm

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

0.5503934621810913 0.29583078622817993

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

0.5503906011581421 0.2958293557167053

Source code in tinygrad/nn/__init__.py

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: Optional[Tensor] = Tensor.ones(sz) if affine else None
  self.bias: Optional[Tensor] = Tensor.zeros(sz) if affine else None

  self.num_batches_tracked = Tensor.zeros(1, dtype='long' if is_dtype_supported(dtypes.long) else 'int', 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: Union[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.6944 0.885  0.0651 0.8116]]]

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

[[[ 0.2012 -0.5826]]]

Source code in tinygrad/nn/__init__.py

def Conv1d(in_channels:int, out_channels:int, kernel_size:int, stride=1, padding:Union[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: Union[int, Tuple[int, ...]],
    stride=1,
    padding: Union[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.429  0.8286 0.9955 0.9372]
   [0.7824 0.5034 0.5371 0.2154]
   [0.7706 0.6966 0.6172 0.9011]
   [0.2546 0.4914 0.7198 0.4352]]]]

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

[[[[-0.5456 -0.5856]
   [-0.422  -0.4364]]]]

Source code in tinygrad/nn/__init__.py

def __init__(self, in_channels:int, out_channels:int, kernel_size:Union[int, Tuple[int, ...]], stride=1, padding:Union[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: Optional[Tensor] = 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.6016 0.7056 0.315  0.83  ]]]

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

[[[ 0.1591 -0.0357  0.0068  0.4761 -0.338   0.245 ]]]

Source code in tinygrad/nn/__init__.py

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: Union[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.4496 0.6639 0.6248 0.0882]
   [0.1408 0.3608 0.0165 0.4955]
   [0.7588 0.3654 0.1834 0.7262]
   [0.3087 0.0625 0.0561 0.2665]]]]

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

[[[[0.2168 0.1076 0.0712 0.1739 0.2734 0.3002]
   [0.3153 0.2433 0.3367 0.2862 0.3344 0.306 ]
   [0.2234 0.1647 0.0879 0.0178 0.0054 0.3538]
   [0.3273 0.2546 0.3318 0.3341 0.1855 0.3134]
   [0.419  0.2406 0.1625 0.3132 0.1665 0.1836]
   [0.3387 0.2726 0.2304 0.3126 0.2574 0.2411]]]]

Source code in tinygrad/nn/__init__.py

def __init__(self, in_channels:int, out_channels:int, kernel_size:Union[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.9726 0.6265 0.8876]
 [0.3015 0.8213 0.1333]]

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

[[ 0.9665  0.9826 -0.5089  0.9712]
 [ 0.5128  0.2275 -0.2947  0.3601]]

Source code in tinygrad/nn/__init__.py

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.

• Described: https://paperswithcode.com/method/group-normalization
• Paper: https://arxiv.org/abs/1803.08494v3

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

2.0031216144561768 0.5758187174797058

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

4.539322162600001e-08 1.0012893676757812

Source code in tinygrad/nn/__init__.py

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: Optional[Tensor] = Tensor.ones(num_channels) if affine else None
  self.bias: Optional[Tensor] = Tensor.zeros(num_channels) if affine else None

InstanceNorm

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

Applies Instance Normalization over a mini-batch of inputs.

• Described: https://paperswithcode.com/method/instance-normalization
• Paper: https://arxiv.org/abs/1607.08022v3

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

2.0593199729919434 0.5625905394554138

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

4.370501116568448e-08 1.0052322149276733

Source code in tinygrad/nn/__init__.py

def __init__(self, num_features:int, eps=1e-5, affine=True):
  self.num_features, self.eps = num_features, eps
  self.weight: Optional[Tensor] = Tensor.ones(num_features) if affine else None
  self.bias: Optional[Tensor] = Tensor.zeros(num_features) if affine else None

LayerNorm

LayerNorm(
    normalized_shape: Union[int, Tuple[int, ...]],
    eps=1e-05,
    elementwise_affine=True,
)

Applies Layer Normalization over a mini-batch of inputs.

• Described: https://paperswithcode.com/method/layer-normalization
• Paper: https://arxiv.org/abs/1607.06450v1

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

2.2973506450653076 0.5354159474372864

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

-4.42120665411494e-07 1.0169609785079956

Source code in tinygrad/nn/__init__.py

def __init__(self, normalized_shape:Union[int, Tuple[int, ...]], eps=1e-5, elementwise_affine=True):
  self.normalized_shape: Tuple[int, ...] = make_tuple(normalized_shape, 1)
  self.axis, self.eps, self.elementwise_affine = tuple(-1-i for i in range(len(self.normalized_shape))), eps, elementwise_affine
  self.weight: Optional[Tensor] = Tensor.ones(*self.normalized_shape) if elementwise_affine else None
  self.bias: Optional[Tensor] = Tensor.zeros(*self.normalized_shape) if elementwise_affine else None

LayerNorm2d

LayerNorm2d(
    normalized_shape: Union[int, Tuple[int, ...]],
    eps=1e-05,
    elementwise_affine=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())

2.0070536136627197 0.5694392919540405

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

-2.434611587887048e-07 1.005044937133789

Source code in tinygrad/nn/__init__.py

def __init__(self, normalized_shape:Union[int, Tuple[int, ...]], eps=1e-5, elementwise_affine=True):
  self.normalized_shape: Tuple[int, ...] = make_tuple(normalized_shape, 1)
  self.axis, self.eps, self.elementwise_affine = tuple(-1-i for i in range(len(self.normalized_shape))), eps, elementwise_affine
  self.weight: Optional[Tensor] = Tensor.ones(*self.normalized_shape) if elementwise_affine else None
  self.bias: Optional[Tensor] = Tensor.zeros(*self.normalized_shape) if elementwise_affine else None

RMSNorm

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

Applies Root Mean Square Normalization to input.

• Described: https://paperswithcode.com/method/rmsnorm
• Paper: https://arxiv.org/abs/1910.07467

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

def __init__(self, dim:int, eps=1e-6): self.eps, self.weight = eps, Tensor.ones(dim)

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.5369 -0.4682  0.1494]
 [-0.0284 -0.102  -0.3805]
 [-0.3976  0.4794 -0.4422]
 [-0.5369 -0.4682  0.1494]]

Source code in tinygrad/nn/__init__.py

def __init__(self, vocab_size:int, embed_size:int):
  self.vocab_sz, self.embed_sz, self.weight = vocab_size, embed_size, 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

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: Optional[Tensor] = Tensor.zeros(hidden_size*4) if bias else None
  self.bias_hh: Optional[Tensor] = 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,
)

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.

• Described: https://paperswithcode.com/method/sgd

Source code in tinygrad/nn/optim.py

def SGD(params: List[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov=False, classic=False):
  """
  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.

  - Described: https://paperswithcode.com/method/sgd
  """
  return LARS(params, lr, momentum, weight_decay, nesterov, classic, tcoef=0.0)

LARS

LARS(
    params: List[Tensor],
    lr=0.001,
    momentum=0.9,
    weight_decay=0.0001,
    nesterov=False,
    classic=True,
    tcoef=0.001,
)

Bases: Optimizer

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

• Described: https://paperswithcode.com/method/lars
• Paper: https://arxiv.org/abs/1708.03888v3

Source code in tinygrad/nn/optim.py

def __init__(self, params:List[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, nesterov=False, classic=True, tcoef=0.001):
  super().__init__(params, lr)
  self.momentum, self.wd, self.nesterov, self.classic, self.tcoef = momentum, weight_decay, nesterov, classic, tcoef
  self.b = [Tensor.zeros(*t.shape, dtype=t.dtype, device=t.device, requires_grad=False) for t in self.params] 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,
)

AdamW optimizer with optional weight decay.

• Described: https://paperswithcode.com/method/adamw
• Paper: https://arxiv.org/abs/1711.05101v3

Source code in tinygrad/nn/optim.py

def AdamW(params: List[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.01):
  """
  AdamW optimizer with optional weight decay.

  - Described: https://paperswithcode.com/method/adamw
  - Paper: https://arxiv.org/abs/1711.05101v3
  """
  return LAMB(params, lr, b1, b2, eps, weight_decay, adam=True)

Adam

Adam(
    params: List[Tensor],
    lr=0.001,
    b1=0.9,
    b2=0.999,
    eps=1e-08,
)

Adam optimizer.

• Described: https://paperswithcode.com/method/adam
• Paper: https://arxiv.org/abs/1412.6980

Source code in tinygrad/nn/optim.py

def Adam(params: List[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8):
  """
  Adam optimizer.

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

LAMB

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

Bases: Optimizer

LAMB optimizer with optional weight decay.

• Described: https://paperswithcode.com/method/lamb
• Paper: https://arxiv.org/abs/1904.00962

Source code in tinygrad/nn/optim.py

def __init__(self, params: List[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, adam=False):
  super().__init__(params, lr)
  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).contiguous() for _ in [b1, b2])
  self.m = [Tensor.zeros(*t.shape, dtype=dtypes.float32, device=t.device, requires_grad=False).contiguous() for t in self.params]
  self.v = [Tensor.zeros(*t.shape, dtype=dtypes.float32, device=t.device, requires_grad=False).contiguous() for t in self.params]

Load/Save

safe_load

safe_load(
    fn: Union[Tensor, str, Path]
) -> Dict[str, Tensor]

Loads a .safetensor file from disk, returning the state_dict.

state_dict = nn.state.safe_load("test.safetensor")

Source code in tinygrad/nn/state.py

def safe_load(fn:Union[Tensor, str, pathlib.Path]) -> Dict[str, Tensor]:
  """
  Loads a .safetensor file from disk, 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: Optional[Dict[str, Any]] = 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

def safe_save(tensors:Dict[str, Tensor], fn:str, metadata:Optional[Dict[str, Any]]=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"*((8-len(j)%8)%8)
  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

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

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,
) -> None

Loads a state_dict into a model.

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

def load_state_dict(model, state_dict:Dict[str, Tensor], strict=True, verbose=True, consume=False) -> None:
  """
  Loads a state_dict into a model.

  ```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
  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"):
    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:
        raise ValueError(f'Shape mismatch in layer `{k}`: Expected shape {v.shape}, but found {state_dict[k].shape} in state dict.')
      if isinstance((mlb:=v.lazydata), MultiLazyBuffer):
        if isinstance(state_dict[k].lazydata, MultiLazyBuffer): v.replace(state_dict[k]).realize()
        else: v.replace(state_dict[k].shard(mlb.device, mlb.axis)).realize()
      else: v.replace(state_dict[k].to(v.device)).realize()
      if consume: del state_dict[k]

torch_load

torch_load(t: Tensor) -> Dict[str, Tensor]

Loads a torch .pth file from disk.

state_dict = nn.state.torch_load("test.pth")

Source code in tinygrad/nn/state.py

@accept_filename
def torch_load(t:Tensor) -> Dict[str, Tensor]:
  """
  Loads a torch .pth file from disk.

  ```python
  state_dict = nn.state.torch_load("test.pth")
  ```
  """
  offsets: Dict[Union[str, int], int] = {}
  lens: Dict[Union[str, int], int] = {}
  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 offsets: return None
    byte_offset = offsets[storage[2]]+storage_offset*storage[1].itemsize
    ret = t[byte_offset:byte_offset+prod(size)*storage[1].itemsize].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. CLANG to permute {intermediate_shape} with {permute_indexes}")
      assert storage[1] != dtypes.bfloat16, "can't CLANG permute BF16"
      # TODO: find a nice way to support all shapetracker 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_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
    myzip = zipfile.ZipFile(fobj, 'r')
    base_name = myzip.namelist()[0].split('/', 1)[0]
    for n in myzip.namelist():
      if n.startswith(f'{base_name}/data/'):
        with myzip.open(n) as myfile:
          offsets[n.split("/")[-1]] = myfile._orig_compress_start # type: ignore
    with myzip.open(f'{base_name}/data.pkl') as myfile:
      return TorchPickle(myfile).load()
  elif passthrough_reset(tarfile.is_tarfile(fobj)): # NOTE: passthrough_reset required to support python < 3.11
    with tarfile.open(fileobj=fobj, mode="r") as tar:
      storages_offset = tar.getmember('storages').offset_data
      f = unwrap(tar.extractfile('storages'))
      for i in range(TorchPickle(f).load()):  # num_storages
        (key, _, storage_type), sz = TorchPickle(f).load(), struct.unpack('<q', f.read(8))[0]
        offsets[key] = storages_offset + f.tell()
        f.seek(sz*storage_type.itemsize, 1)
      f = unwrap(tar.extractfile('tensors'))
      for _ in range(TorchPickle(f).load()):  # num_tensors
        (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)
      return {k:v.tensor if isinstance(v, Parameter) else v for k,v in TorchPickle(unwrap(tar.extractfile('pickle'))).load().items()}
  else:
    pkl = TorchPickle(fobj)
    _, _, _, rwd, _, ids, base_offset = pkl.load(), pkl.load(), pkl.load(), fobj.tell(), pkl.load(), pkl.load(), fobj.tell()
    for i in ids:
      offsets[i] = base_offset + 8
      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 from a tensor.

fn = "Meta-Llama-3-8B-Instruct.Q4_0.gguf"
gguf_tensor = Tensor.empty(os.stat(fn).st_size, dtype=dtypes.uint8, device=f"disk:{fn}").to(Device.DEFAULT)
kv_data, state_dict = gguf_load(gguf_tensor)

Source code in tinygrad/nn/state.py

@accept_filename
def gguf_load(tensor: Tensor) -> Tuple[Dict, Dict[str, Tensor]]:
  """
  Loads a gguf file from a tensor.

  ```python
  fn = "Meta-Llama-3-8B-Instruct.Q4_0.gguf"
  gguf_tensor = Tensor.empty(os.stat(fn).st_size, dtype=dtypes.uint8, device=f"disk:{fn}").to(Device.DEFAULT)
  kv_data, state_dict = gguf_load(gguf_tensor)
  ```
  """
  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