创建

@staticmethod
def empty(*shape, device:str|tuple[str, ...]|None=None, dtype:DTypeLike|None=None, **kwargs) -> Tensor:
  """
  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)
  ```
  """
  dtype, shape = to_dtype(dtype) if dtype is not None else dtypes.default_float, argfix(*shape)
  if not isinstance(size:=prod([x.vmax if isinstance(x, UOp) else x for x in shape]), int): raise ValueError(f"size must be int {size}")
  if isinstance(device, tuple):
    return Tensor(UOp.multi(*[UOp.new_buffer(Device.canonicalize(d), size, dtype).reshape(shape) for d in device], axis=None),
                  device, dtype, **kwargs)
  return Tensor(UOp.new_buffer(Device.canonicalize(device), size, dtype), device, dtype, **kwargs).reshape(shape)

@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)

@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)

@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)

@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
  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 (output_len:=ceildiv(stop-start, step)) <= 0: return Tensor([], dtype=dtype, **kwargs)
  return (Tensor.full((output_len,), step, dtype=dtype, **kwargs)._cumalu(0, Ops.ADD) + (start - step)).cast(dtype)

@staticmethod
def linspace(start:int|float, stop:int|float, steps:int, **kwargs) -> Tensor:
  """
  Returns a 1-D tensor of `steps` evenly spaced values from `start` to `stop`, inclusive.

  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.linspace(0, 10, 5).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.linspace(-1, 1, 5).numpy())
  ```
  """
  if steps < 0: raise ValueError("number of steps must be non-negative")
  if (dtype := to_dtype(kwargs.pop("dtype", dtypes.default_float))) == dtypes.bool: raise ValueError("linspace with bool dtype is not supported")
  if steps == 1: return Tensor([start], dtype=dtype, **kwargs)
  return (start + Tensor.arange(steps, **kwargs) * ((stop - start) / (steps - 1))).cast(dtype)

@staticmethod
def eye(n:int, m:int|None=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)))

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)

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)

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)

@staticmethod
def from_blob(ptr:int, shape:tuple[int, ...], **kwargs) -> Tensor:
  """
  Exposes the pointer as a Tensor without taking ownership of the original data.
  The pointer must remain valid for the entire lifetime of the created Tensor.

  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.
  """
  r = Tensor.empty(*shape, **kwargs)
  r.lazydata.buffer.allocate(external_ptr=ptr)
  return r

@staticmethod
def from_url(url:str, gunzip:bool=False, **kwargs) -> Tensor:
  """
  Create a Tensor from a URL.

  This is the preferred way to access Internet resources.
  It currently returns a DISK Tensor, but in the future it may return an HTTP Tensor.
  This also will soon become lazy (when possible) and not print progress without DEBUG.

  THe `gunzip` flag will gzip extract the resource and return an extracted Tensor.
  """
  return Tensor(fetch(url, gunzip=gunzip), **kwargs)

@staticmethod
def manual_seed(seed=0) -> None:
  """
  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._device_seeds, Tensor._device_rng_counters = seed, {}, {}

@staticmethod
def rand(*shape, device:str|None=None, dtype:DTypeLike|None=None, contiguous:bool=True, **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 device is not None and not isinstance(device, str): raise ValueError(f"rand only supports single device, got {device=}")
  device = Device.canonicalize(device)

  # if shape has 0, return zero tensor
  if (numel := prod(shape)) == 0: return Tensor.zeros(shape, device=device, dtype=dtype, **kwargs)
  num = ceildiv(numel * dtype.itemsize, 4)

  # generate per device seeds and rng counter if we haven't seen this device yet
  if device not in Tensor._device_seeds:
    Tensor._device_seeds[device] = Tensor(
      [int.from_bytes(hashlib.sha256(len(Tensor._device_seeds).to_bytes(4, "big")).digest(), "big"), Tensor._seed],
      device=device, dtype=dtypes.uint32, requires_grad=False)
    Tensor._device_rng_counters[device] = Tensor([0], device=device, dtype=dtypes.uint32, requires_grad=False)
  # increment rng counter for devices
  else: Tensor._device_rng_counters[device].assign(Tensor._device_rng_counters[device] + num).contiguous()

  # threefry random bits
  counts0 = (Tensor.arange(ceildiv(num, 2), device=device, dtype=dtypes.uint32, requires_grad=False)+Tensor._device_rng_counters[device])
  counts1 = counts0 + ceildiv(num, 2)
  bits = Tensor._threefry_random_bits(Tensor._device_seeds[device], counts0, counts1)[: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 and reshape
  out = bits.bitcast(dtype)[:numel].sub(1).reshape(shape).requires_grad_(kwargs.get("requires_grad"))
  return out.contiguous() if contiguous else out

def rand_like(self, **kwargs) -> Tensor:
  """
  Creates a tensor with the same shape and sharding as `self`, 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"
  t = Tensor.ones(2, 3)
  print(Tensor.rand_like(t).numpy())
  ```
  """
  dtype = kwargs.pop("dtype", self.dtype)
  if isinstance(self.device, tuple):
    if kwargs.get("device") is not None: raise RuntimeError("cannot specify `device` on `rand_like` of a multi device tensor")
    if self.lazydata.axis is None: return Tensor.rand(*self.shape, dtype=dtype, **kwargs).shard(self.device)
    contiguous = kwargs.pop("contiguous", True)
    sharded_shape = tuple(s//len(self.device) if a==self.lazydata.axis else s for a,s in enumerate(self.shape))
    rands = [Tensor.rand(sharded_shape, device=d, dtype=dtype, contiguous=contiguous, **kwargs).lazydata for d in self.device]
    return Tensor(UOp.multi(*rands, axis=self.lazydata.axis), device=self.device, dtype=dtype, **kwargs)
  return Tensor.rand(*self.shape, device=kwargs.pop("device", self.device), dtype=dtype, **kwargs)

@staticmethod
def randn(*shape, dtype:DTypeLike|None=None, requires_grad:bool|None=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)).requires_grad_(requires_grad)

@staticmethod
def randint(*shape, low=0, high=10, dtype=dtypes.int32, **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 = to_dtype(dtype)
  if not dtypes.is_int(dtype): raise TypeError(f"{dtype=} must be int")
  return Tensor.uniform(*shape, low=low, high=high, dtype=dtype, **kwargs)

@staticmethod
def normal(*shape, mean=0.0, std=1.0, requires_grad:bool|None=None, **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).requires_grad_(requires_grad)

@staticmethod
def uniform(*shape, low=0.0, high=1.0, dtype:DTypeLike|None=None, requires_grad:bool|None=None, **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())
  ```
  """
  return (((high-low) * Tensor.rand(*shape, **kwargs)).cast(dtype or dtypes.default_float) + low).requires_grad_(requires_grad)

@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)

@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)

@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)

@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)