属性

def numel(self) -> sint:
  """
  Returns the total number of elements in the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
  print(t.numel())
  ```
  """
  return prod(self.shape)

def element_size(self) -> int:
  """
  Returns the size in bytes of an individual element in the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([5], dtype=dtypes.int16)
  print(t.element_size())
  ```
  """
  return self.dtype.itemsize

def nbytes(self) -> int:
  """
  Returns the total number of bytes of all elements in the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([8, 9], dtype=dtypes.float)
  print(t.nbytes())
  ```
  """
  return self.numel() * self.element_size()

def is_floating_point(self) -> bool:
  """
  Returns `True` if the tensor contains floating point types, i.e. is one of `dtype.float64`, `dtype.float32`,
  `dtype.float16`, `dtype.bfloat16`.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([8, 9], dtype=dtypes.float32)
  print(t.is_floating_point())
  ```
  """
  return dtypes.is_float(self.dtype)

def size(self, dim:int|None=None) -> sint|tuple[sint, ...]:
  """
  Return the size of the tensor. If `dim` is specified, return the length along dimension `dim`. Otherwise return the shape of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[4, 5, 6], [7, 8, 9]])
  print(t.size())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.size(dim=1))
  ```
  """
  return self.shape if dim is None else self.shape[dim]

def data(self) -> memoryview:
  """
  Returns the data of this tensor as a memoryview.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1, 2, 3, 4])
  print(np.frombuffer(t.data(), dtype=np.int32))
  ```
  """
  if 0 in self.shape: return memoryview(bytearray(0)).cast(self.dtype.base.fmt)
  assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
  return self._buffer().as_typed_buffer(self.shape)

def item(self) -> ConstType:
  """
  Returns the value of this tensor as a standard Python number.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor(42)
  print(t.item())
  ```
  """
  assert self.numel() == 1, "must have one element for item"
  return self.data()[(0,) * len(self.shape)]

def tolist(self) -> Sequence[ConstType]|ConstType:
  """
  Returns the value of this tensor as a nested list.
  Returns single value for const tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1, 2, 3, 4])
  print(t.tolist())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor(5)
  print(t.tolist())
  ```
  """
  return self.data().tolist()

def numpy(self) -> 'np.ndarray':  # type: ignore [name-defined] # noqa: F821
  """
  Returns the value of this tensor as a `numpy.ndarray`.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1, 2, 3, 4])
  print(repr(t.numpy()))
  ```
  """
  assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
  import numpy as np
  if self.dtype.base == dtypes.bfloat16: return self.float().numpy()
  if 0 in self.shape: return np.empty(self.shape, dtype=_to_np_dtype(self.dtype.base))
  return self._buffer().numpy().reshape(self.shape)

def schedule_with_vars(self, *lst:Tensor) -> tuple[list[ScheduleItem], dict[Variable, int]]:
  """
  Creates the schedule needed to realize these Tensor(s), with Variables.

  NOTE: A Tensor can only be scheduled once.
  """
  big_sink = UOp.sink(*[x.lazydata for x in (self,)+lst])

  # TODO: move this to scheduler tensor_map pass
  if any(x.op is Ops.MULTI for x in big_sink.toposort):
    # multi fixup
    _apply_map_to_tensors(get_multi_map(big_sink), name="Apply Multi Map")
    big_sink = UOp.sink(*flatten([x.lazydata.src if x.lazydata.op is Ops.MULTI else [x.lazydata] for x in (self,)+lst]))

  # verify Tensors match the spec
  if __debug__: type_verify(list(big_sink.toposort), tensor_uop_spec)

  schedule, var_vals, becomes_map = create_schedule_with_vars(big_sink)
  _apply_map_to_tensors(becomes_map, name="Apply Schedule Map")
  return memory_planner(schedule), var_vals

def schedule(self, *lst:Tensor) -> list[ScheduleItem]:
  """Creates the schedule needed to realize these Tensor(s)."""
  schedule, var_vals = self.schedule_with_vars(*lst)
  assert len(var_vals) == 0
  return schedule

def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
  """Triggers the computation needed to create these Tensor(s)."""
  run_schedule(*self.schedule_with_vars(*lst), do_update_stats=do_update_stats)
  return self

def replace(self, x:Tensor, allow_shape_mismatch=False) -> Tensor:
  """
  Replaces the data of this tensor with the data of another tensor. Only the shape of the tensors must match.
  """
  # used for replacing a Tensor with a new version of it (potentially with a different device and dtype)
  assert self.shape == x.shape or allow_shape_mismatch, f"replace shape mismatch {self.shape} != {x.shape}"
  self.lazydata = x.lazydata
  return self

def assign(self, x) -> Tensor:
  # TODO: this is a hack for writing to DISK. remove with working assign
  if isinstance(self.device, str) and self.device.startswith("DISK"):
    if x.__class__ is not Tensor: x = Tensor(x, device="CPU", dtype=self.dtype)
    self.contiguous().realize().lazydata.base.buffer.ensure_allocated().copyin(x._data())
    return self
  if x.__class__ is not Tensor: x = Tensor(x, device=self.device, dtype=self.dtype)
  if self.lazydata is x.lazydata: return self  # a self assign is a NOOP
  # NOTE: we allow cross device assign
  assert self.shape == x.shape, f"assign shape mismatch {self.shape} != {x.shape}"
  assert self.device == x.device, f"assign device mismatch {self.device} != {x.device}"
  assert self.dtype == x.dtype, f"assign dtype mismatch {self.dtype} != {x.dtype}"
  assert not x.requires_grad  # self requires_grad is okay?
  if not self.lazydata.is_realized: return self.replace(x)
  self.lazydata = self.lazydata.assign(x.lazydata)
  return self

def detach(self) -> Tensor:
  """
  Returns a new tensor with the same data as this tensor, but detached from the autograd graph.
  """
  return Tensor(self.lazydata.detach(), device=self.device, requires_grad=False)

def clone(self) -> Tensor:
  """
  Creates a clone of this tensor allocating a separate buffer for the data.
  """
  ret = Tensor(self.lazydata.clone(), self.device, requires_grad=self.requires_grad)
  if self.grad is not None: ret.grad = self.grad.clone()
  return ret

def to(self, device:str|tuple[str, ...]|None) -> Tensor:
  """
  Moves the tensor to the given device.
  """
  device = tuple(Device.canonicalize(x) for x in device) if isinstance(device, (tuple, list)) else Device.canonicalize(device)
  if device == self.device: return self
  if not isinstance(device, str): return self.shard(device)
  ret = Tensor(self.lazydata, device, requires_grad=self.requires_grad)
  if self.grad is not None: ret.grad = self.grad.to(device)
  return ret

def to_(self, device:str|tuple[str, ...]|None) -> Tensor:
  """
  Moves the tensor to the given device in place.
  """
  real = self.to(device)
  if self.grad is not None and real.grad is not None: self.grad.replace(real.grad)
  return self.replace(real)

def shard(self, devices:tuple[str, ...], axis:int|None=None) -> Tensor:
  """
  Shards the tensor across the given devices. Optionally specify which axis to shard on.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.empty(2, 4)
  print(t.shard((t.device, t.device), axis=1).lazydata)
  ```
  """
  assert isinstance(self.device, str), "can't shard a MultiLazyBuffer"
  devices = tuple(Device.canonicalize(x) for x in devices)
  mlb = self.lazydata.shard(devices, self._resolve_dim(axis) if axis is not None else None)
  return Tensor(mlb, device=devices, requires_grad=self.requires_grad)

def shard_(self, devices:tuple[str, ...], axis:int|None=None) -> Tensor:
  """
  Shards the tensor across the given devices in place.
  """
  return self.replace(self.shard(devices, axis))

def contiguous(self) -> Tensor:
  """
  Returns a contiguous tensor.
  """
  return self._apply_uop(UOp.contiguous)

def contiguous_backward(self) -> Tensor:
  """
  Inserts a contiguous operation in the backward pass.
  """
  return self._apply_uop(UOp.contiguous_backward)

def gradient(self, *targets:Tensor, gradient:Tensor|None=None, materialize_grads=False) -> list[Tensor]:
  """
  Compute the gradient of the targets with respect to self.

  ```python exec="true" source="above" session="tensor" result="python"
  x = Tensor.eye(3)
  y = Tensor([[2.0,0,-2.0]])
  z = y.matmul(x).sum()
  dx, dy = z.gradient(x, y)

  print(dx.tolist())  # dz/dx
  print(dy.tolist())  # dz/dy
  ```
  """
  assert gradient is not None or self.shape == tuple(), "when no gradient is provided, backward must be called on a scalar tensor"
  if gradient is None: gradient = Tensor(1.0, dtype=self.dtype, device=self.device, requires_grad=False)
  rets = []
  target_uops = [x.lazydata for x in targets]
  grads = compute_gradient(self.lazydata, gradient.lazydata, set(target_uops))
  ret = []
  for x in target_uops:
    if (y:=grads.get(x)) is None:
      if materialize_grads: y = x.const_like(0)
      else: raise RuntimeError(f"{x}\n\nnot found in\n\n{self.lazydata}")
    ret.append(y)
  rets.append(ret)
  # create returned Tensors
  return [Tensor(u, device=t.device) for t,u in zip(targets, rets[0])]

def backward(self, gradient:Tensor|None=None) -> Tensor:
  """
  Propagates the gradient of a tensor backwards through the computation graph.
  If the 'gradient' argument is not provided, the tensor must be a scalar, and the gradient is implicitly set to 1.0.
  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
  t.sum().backward()
  print(t.grad.numpy())
  ```
  """
  all_uops = self.lazydata.toposort
  tensors_need_grad: list[Tensor] = [t for tref in all_tensors if (t:=tref()) is not None and \
                                     t.lazydata in all_uops and t.requires_grad and not Tensor.no_grad]
  # clear contexts
  for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient, materialize_grads=True)):
    assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
    t.grad = g if t.grad is None else (t.grad + g)
  return self