复杂运算

def sum(self, axis:int|Sequence[int]|None=None, keepdim=False, dtype:DTypeLike|None=None) -> Tensor:
  """
  Returns the sum of the elements of the tensor along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the maximum is computed and whether the reduced dimensions are retained.

  You can pass in `dtype` keyword argument to control the data type of the accumulation.
  If not specified, the accumulation data type is chosen based on the input tensor's data type.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(6).reshape(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.sum().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.sum(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.sum(axis=1).numpy())
  ```
  """
  ret = self.cast(sum_acc_dtype(self.dtype) if dtype is None else dtype)._reduce(Ops.ADD, axis, keepdim)
  return ret.cast(self.dtype) if dtype is None and self.dtype in (dtypes.float16, dtypes.bfloat16) else ret

def prod(self, axis:int|Sequence[int]|None=None, keepdim=False, dtype:DTypeLike|None=None) -> Tensor:
  """
  Returns the product of the elements of the tensor along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the maximum is computed and whether the reduced dimensions are retained.

  You can pass in `dtype` keyword argument to control the data type of the accumulation.
  If not specified, the accumulation data type is chosen based on the input tensor's data type.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([-1, -2, -3, 1, 2, 3]).reshape(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.prod().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.prod(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.prod(axis=1).numpy())
  ```
  """
  return self.cast(dtype if dtype is not None else self.dtype)._reduce(Ops.MUL, axis, keepdim)

def max(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor:
  """
  Returns the maximum value of the tensor along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the maximum is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 0, 2], [5, 4, 3]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.max().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.max(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.max(axis=1, keepdim=True).numpy())
  ```
  """
  return self._reduce(Ops.MAX, axis, keepdim)

def min(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor:
  """
  Returns the minimum value of the tensor along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the minimum is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 0, 2], [5, 4, 3]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.min().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.min(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.min(axis=1, keepdim=True).numpy())
  ```
  """
  return self._inverse().max(axis=axis, keepdim=keepdim)._inverse()

def any(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor:
  """
  Tests if any element evaluates to `True` along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the reduce axis and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[True, True], [True, False], [False, False]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.any().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.any(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.any(axis=1, keepdim=True).numpy())
  ```
  """
  return self.bool().max(axis, keepdim)

def all(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor:
  """
  Tests if all element evaluates to `True` along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the reduce axis and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[True, True], [True, False], [False, False]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.all().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.all(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.all(axis=1, keepdim=True).numpy())
  ```
  """
  return self.logical_not().any(axis, keepdim).logical_not()

def isclose(self, other:Tensor, rtol:float=1e-05, atol:float=1e-08, equal_nan=False) -> Tensor:
  """
  Returns a new tensor with element-wise comparison of closeness to `other` within a tolerance.

  The `rtol` and `atol` keyword arguments control the relative and absolute tolerance of the comparison.

  By default, two `NaN` values are not close to each other. If `equal_nan` is `True`, two `NaN` values are considered close.

  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor([1e-7, 1e-8, 1e-9, float('nan')]).isclose(Tensor([0.0, 0.0, 0.0, float('nan')])).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor([float('nan')]).isclose(Tensor([float('nan')]), equal_nan=True).numpy())
  ```
  """
  is_finite_close = self.isfinite() & other.isfinite() & ((self - other).abs() <= atol + rtol * other.abs())
  is_infinite_close = (self.isinf() | other.isinf()) & (self == other)
  is_nan_close = (self.isnan() & other.isnan()) & equal_nan
  return is_finite_close | is_infinite_close | is_nan_close

def mean(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor:
  """
  Returns the mean value of the tensor along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the mean is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.normal(2, 3, mean=2.5, std=0.5)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.mean().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.mean(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.mean(axis=1).numpy())
  ```
  """
  output_dtype = self.dtype if dtypes.is_float(self.dtype) else dtypes.float32
  numerator = self.cast(sum_acc_dtype(self.dtype)).sum(axis=axis, keepdim=keepdim)
  return numerator.div(prod([cast(int, si) for si, so in zip(self.shape, self.sum(axis=axis, keepdim=True).shape) if resolve(si != so)])) \
    .cast(output_dtype)

def var(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> Tensor:
  """
  Returns the variance of the tensor along the specified axis or axes.

  You can pass in `axis`, `keepdim`, and `correction` keyword arguments to control the axis along
  which the variance is computed, whether the reduced dimensions are retained, and the Bessel's correction applied.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.normal(2, 3, mean=2.5, std=0.5)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.var().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.var(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.var(axis=1).numpy())
  ```
  """
  squares = (self - self.mean(axis=axis, keepdim=True)).square()
  n = prod([si for si, so in zip(self.shape, squares.sum(axis=axis, keepdim=True).shape) if resolve(si != so)])
  return squares.sum(axis=axis, keepdim=keepdim).div(smax([0, n-correction]))

def var_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Tensor, Tensor]:
  """
  Calculates the variance and mean over the dimensions specified by dim.
  Syntactic sugar around `Tensor.var` and `Tensor.mean` to match `torch.var_mean`.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.normal(2, 3, mean=2.5, std=0.5)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  var, mean = t.var_mean()
  print(var.numpy(), mean.numpy())
  ```
  """
  return self.var(axis, keepdim, correction), self.mean(axis, keepdim)

def std(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> Tensor:
  """
  Returns the standard deviation of the tensor along the specified axis or axes.

  You can pass in `axis`, `keepdim`, and `correction` keyword arguments to control the axis along
  which the standard deviation is computed, whether the reduced dimensions are retained, and the Bessel's correction applied.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.normal(2, 3, mean=2.5, std=0.5)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.std().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.std(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.std(axis=1).numpy())
  ```
  """
  return self.var(axis, keepdim, correction).sqrt()

def std_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Tensor, Tensor]:
  """
  Calculates the standard deviation and mean over the dimensions specified by dim.
  Syntactic sugar around `Tensor.std` and `Tensor.mean` to match `torch.std_mean`.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.normal(2, 3, mean=2.5, std=0.5)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  std, mean = t.std_mean()
  print(std.numpy(), mean.numpy())
  ```
  """
  return self.std(axis, keepdim, correction), self.mean(axis, keepdim)

def softmax(self, axis=-1, dtype:DTypeLike|None=None) -> Tensor:
  """
  Applies the softmax function to the tensor along the specified axis.

  Rescales the elements of the tensor such that they lie in the range [0, 1] and sum to 1.

  You can pass in the `axis` keyword argument to control the axis along which the softmax is computed.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.randn(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.softmax().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.softmax(axis=0).numpy())
  ```
  """
  if getenv("SINGLE_KERNEL_SOFTMAX"):
    _, e, ss = self.contiguous()._softmax(axis, dtype)
    return e.div(ss).fuse()
  _, e, ss = self._softmax(axis, dtype)
  return e.div(ss)

def log_softmax(self, axis=-1, dtype:DTypeLike|None=None) -> Tensor:
  """
  Applies the log-softmax function to the tensor along the specified axis.

  The log-softmax function is a numerically stable alternative to the softmax function in log space.

  You can pass in the `axis` keyword argument to control the axis along which the log-softmax is computed.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.randn(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.log_softmax().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.log_softmax(axis=0).numpy())
  ```
  """
  m, _, ss = self._softmax(axis, dtype)
  return m - ss.log()

def logsumexp(self, axis=None, keepdim=False) -> Tensor:
  """
  Computes the log-sum-exp of the tensor along the specified axis or axes.

  The log-sum-exp function is a numerically stable way to compute the logarithm of the sum of exponentials.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the log-sum-exp is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.randn(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logsumexp().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logsumexp(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logsumexp(axis=1).numpy())
  ```
  """
  m = self.max(axis=axis, keepdim=True)
  return (self - m).exp().sum(axis=axis, keepdim=keepdim).log() + m.squeeze(axis)

def logcumsumexp(self, axis=0) -> Tensor:
  """
  Computes the log-cumsum-exp of the tensor along the specified axis or axes.

  The log-cumsum-exp function is a numerically stable way to compute the logarithm of the cumulative sum of exponentials.

  You can pass in the `axis` keyword argument to control the axis along which
  the log-cum-sum-exp is computed.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.randn(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logcumsumexp().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logcumsumexp(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logcumsumexp(axis=1).numpy())
  ```
  """
  if self.ndim == 0: return self
  axis = self._resolve_dim(axis)
  x = self.transpose(axis, -1)
  last_dim_size = x.shape[-1]
  x_reshaped = x.reshape(-1, last_dim_size)
  x_cummax = x_reshaped.cummax(-1).unsqueeze(-1)
  x_expand = x_reshaped.unsqueeze(1).expand(*x_reshaped.shape, last_dim_size)
  mask = Tensor.ones(last_dim_size, last_dim_size, requires_grad=False, device=self.device).tril().unsqueeze(0)
  ret = ((x_expand - x_cummax).exp() * mask).sum(-1).log() + x_cummax.squeeze(-1)
  return ret.reshape(*x.shape).transpose(-1, axis)

def argmax(self, axis=None, keepdim=False) -> Tensor:
  """
  Returns the indices of the maximum value of the tensor along the specified axis.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the maximum is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 0, 2], [5, 4, 3]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmax().numpy()) # Returns the index of the maximum value in the flattened tensor.
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmax(axis=0).numpy()) # Returns the indices of the maximum values along axis 0.
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmax(axis=1).numpy()) # Returns the indices of the maximum values along axis 1.
  ```
  """
  if axis is None: return self.flatten().argmax(0)
  axis = self._resolve_dim(axis)
  m = self == self.max(axis=axis, keepdim=True)
  idx = m * Tensor.arange(self.shape[axis],0,-1, requires_grad=False, device=self.device).reshape(self.shape[axis], *[1]*(self.ndim-axis-1))
  return (self.shape[axis]-idx.max(axis=axis, keepdim=keepdim)).cast(dtypes.int32)

def argmin(self, axis=None, keepdim=False) -> Tensor:
  """
  Returns the indices of the minimum value of the tensor along the specified axis.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the minimum is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 0, 2], [5, 4, 3]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmin().numpy()) # Returns the index of the minimum value in the flattened tensor.
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmin(axis=0).numpy()) # Returns the indices of the minimum values along axis 0.
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmin(axis=1).numpy()) # Returns the indices of the minimum values along axis 1.
  ```
  """
  return self._inverse().argmax(axis=axis, keepdim=keepdim)

def avg_pool2d(self, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0,
               ceil_mode=False, count_include_pad=True) -> Tensor:
  """
  Applies average pooling over a tensor.

  This function supports three different types of `padding`

  1. `int` (single value):
    Applies the same padding value uniformly to all spatial dimensions.

  2. `tuple[int, ...]` (length = number of spatial dimensions):
    Specifies a distinct padding value for each spatial dimension in the form `(padding_height, padding_width, ...)`.

  3. `tuple[int, ...]` (length = 2 * number of spatial dimensions):
    Specifies explicit padding for each side of each spatial dimension in the form
    `(padding_left, padding_right, padding_top, padding_bottom, ...)`.

  When `ceil_mode` is set to `True`, output shape will be determined using ceil division.
  When `count_include_pad` is set to `False`, zero padding will not be included in the averaging calculation.

  NOTE: unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.

  See: https://paperswithcode.com/method/average-pooling

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(25).reshape(1, 1, 5, 5)
  print(t.avg_pool2d().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.avg_pool2d(ceil_mode=True).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.avg_pool2d(padding=1).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.avg_pool2d(padding=1, count_include_pad=False).numpy())
  ```
  """
  axis = tuple(range(-len(k_ := make_tuple(kernel_size, 2)), 0))
  def pool(x:Tensor, padding_:Sequence[int]) -> Tensor: return x.pad(padding_)._pool(k_, stride if stride is not None else k_, dilation)
  reg_pads = self._resolve_pool_pads(padding, len(k_))
  ceil_pads = self._apply_ceil_mode(reg_pads, k_, stride if stride is not None else k_, dilation)
  if not count_include_pad:
    pads = ceil_pads if ceil_mode else reg_pads
    return pool(self, pads).sum(axis) / pool(self.ones_like(), pads).sum(axis)
  if not ceil_mode: return pool(self, reg_pads).mean(axis)
  return pool(self, ceil_pads).sum(axis) / pool(self.pad(reg_pads).ones_like(), tuple(cp-rp for cp,rp in zip(ceil_pads, reg_pads))).sum(axis)

def max_pool2d(self, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0,
               ceil_mode=False, return_indices=False) -> Tensor | tuple[Tensor, Tensor]:
  """
  Applies max pooling over a tensor.

  This function supports three different types of `padding`

  1. `int` (single value):
    Applies the same padding value uniformly to all spatial dimensions.

  2. `tuple[int, ...]` (length = number of spatial dimensions):
    Specifies a distinct padding value for each spatial dimension in the form `(padding_height, padding_width, ...)`.

  3. `tuple[int, ...]` (length = 2 * number of spatial dimensions):
    Specifies explicit padding for each side of each spatial dimension in the form
    `(padding_left, padding_right, padding_top, padding_bottom, ...)`.

  When `ceil_mode` is set to `True`, output shape will be determined using ceil division.
  When `return_indices` is set to `True`, the argmax will be returned along with the max values.

  NOTE: unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.

  See: https://paperswithcode.com/method/max-pooling

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(25).reshape(1, 1, 5, 5)
  print(t.max_pool2d().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.max_pool2d(ceil_mode=True).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.max_pool2d(padding=1).numpy())
  ```
  """
  axis = tuple(range(-len(k_ := make_tuple(kernel_size, 2)), 0))
  pads = self._resolve_pool_pads(padding, len(k_))
  if ceil_mode: pads = self._apply_ceil_mode(pads, k_, stride if stride is not None else k_, dilation)
  pooled = self.pad(pads, value=dtypes.min(self.dtype))._pool(k_, stride if stride is not None else k_, dilation)
  if not return_indices: return pooled.max(axis)
  spatial_sz = math.prod(spatial_shape := self.shape[-len(k_):])
  idx = Tensor.arange(spatial_sz,0,-1, requires_grad=False, device=self.device).reshape(spatial_shape)
  m = pooled == pooled.max(axis, keepdim=True)
  idx = m * idx.pad(pads, value=dtypes.min(idx.dtype))._pool(k_, stride if stride is not None else k_, dilation)
  return pooled.max(axis), spatial_sz - idx.max(axis)

def max_unpool2d(self, indices:Tensor, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0, output_size=None):
  """
  Performs a partial inverse of `max_pool2d` using the indices from the argmax.

  When `output_size` is provided, the output shape disambiguates to the provided shape.

  NOTE: unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(1, 17).reshape(1, 1, 4, 4)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  output, indices = Tensor.max_pool2d(t, return_indices=True)
  print(output.numpy())
  print(indices.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.max_unpool2d(output, indices).numpy())
  ```
  """
  bs,c,*spatial_shape = self.shape
  if output_size is None:
    k_,d_,s_ = (make_tuple(x, len(spatial_shape)) for x in (kernel_size, dilation, stride if stride is not None else kernel_size))
    p_ = _flat_to_grouped(self._resolve_pool_pads(padding, len(spatial_shape)))
    # https://arxiv.org/pdf/1603.07285 inverse of relationship 15 in section 5.1.
    output_size = tuple((i-1)*s - (pB+pA) + (d*(k-1)+1) for i,k,d,s,(pA,pB) in zip(spatial_shape,k_,d_,s_,p_))
  else: output_size = output_size[-len(spatial_shape):]
  ret = (indices.reshape(bs,c,1,-1)._one_hot_along_dim(prod(output_size), 2) * self.reshape(bs,c,1,-1)).sum(3)
  return ret.reshape(bs,c,*output_size)

def conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding:int|tuple[int, ...]=0,
           dtype:DTypeLike|None=None) -> Tensor:
  """
  Applies a convolution over a tensor with a given `weight` and optional `bias`.

  This function supports three different types of `padding`

  1. `int` (single value):
    Applies the same padding value uniformly to all spatial dimensions.

  2. `tuple[int, ...]` (length = number of spatial dimensions):
    Specifies a distinct padding value for each spatial dimension in the form `(padding_height, padding_width, ...)`.

  3. `tuple[int, ...]` (length = 2 * number of spatial dimensions):
    Specifies explicit padding for each side of each spatial dimension in the form
    `(padding_left, padding_right, padding_top, padding_bottom, ...)`.

  NOTE: unlike PyTorch, this implementation is not limited to only 2d convolutions and instead works for any number of dimensions.

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

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(9).reshape(1, 1, 3, 3)
  w = Tensor.ones(1, 1, 2, 2)
  print(t.conv2d(w).numpy())
  ```
  """
  if IMAGE: return self.image_conv2d(weight, bias, groups, stride, dilation, padding, dtype)
  (bs,cin_), (cout,cin), HW = self.shape[:2], weight.shape[:2], weight.shape[2:]
  padding_ = self._resolve_pool_pads(padding, len(HW))
  assert groups*cin == cin_ and len(self.shape) == len(weight.shape), f"Input Tensor shape {self.shape} does not match the shape of the weights {weight.shape}. ({groups*cin} vs. {cin_})"  # noqa: E501

  # conv2d is a pooling op (with padding)
  x = self.pad(padding_)._pool(HW, stride, dilation)   # (bs, groups*cin, oy, ox, H, W)
  rcout, oyx = cout//groups, x.shape[2:-len(HW)]
  if not all(x == 3 for x in HW) or stride != 1 or dilation != 1 or not WINO:
    # normal conv
    x = x.reshape(bs, groups, cin, 1, *oyx, *HW).expand(bs, groups, cin, rcout, *oyx, *HW).permute(0,1,3,*[4+i for i in range(len(oyx))],2,*[4+len(oyx)+i for i in range(len(HW))])  # noqa: E501

    # conv! broadcasted to (bs, groups, rcout, *oyx, cin, *HW)
    ret = (x * weight.reshape(1, groups, rcout, *[1] * len(oyx), cin, *HW)).sum([-1-i for i in range(1+len(oyx))], keepdim=True, dtype=dtype).reshape(bs, cout, *oyx)  # noqa: E501
    return ret if bias is None else ret.add(bias.reshape(1, -1, *[1] * len(HW)))

  HWI, HWO = (6,) * len(HW), (4,) * len(HW)  # F(4x4,3x3) winograd tiles
  winograd_G = [[1/4, 0, 0], [-1/6, -1/6, -1/6], [-1/6, 1/6, -1/6], [1/24, 1/12, 1/6], [1/24, -1/12, 1/6], [0, 0, 1]]
  winograd_Bt = [[4, 0, -5, 0, 1, 0], [0, -4, -4, 1, 1, 0], [0, 4, -4, -1, 1, 0], [0, -2, -1, 2, 1, 0], [0, 2, -1, -2, 1, 0], [0, 4, 0, -5, 0, 1]]
  winograd_At = [[1, 1, 1, 1, 1, 0], [0, 1, -1, 2, -2, 0], [0, 1, 1, 4, 4, 0], [0, 1, -1, 8, -8, 1]] # applying At in pre-order doubles compile time

  # todo: stride == dilation
  # use padding to round up to 4x4 output tiles
  # (bs, cin_, tyx, HWI)
  d = self.pad(sum([[padding_[i*2], padding_[i*2+1] + (-(dim + sum(padding_[i * 2:(i + 1) * 2]) - 2) % 4)] for i, dim in enumerate(self.shape[-len(HW):])], []))._pool(HWI, HWO)  # noqa: E501
  # move HW to the front: # (HWI, bs, cin_, tyx)
  d = d.permute(*range(len(d.shape)-len(HW),len(d.shape)), *range(len(d.shape)-len(HW)))
  tyx = d.shape[-len(HWI):]  # dim of tiling

  g = weight.permute(*range(len(weight.shape)-len(HW),len(weight.shape)), *range(len(weight.shape)-len(HW)))  # move HW to the front

  # compute 6x6 winograd tiles: GgGt, BtdB
  # (HWI, groups * rcout, cin) -> (HWI, bs=1, groups, rcout, cin, tyx=(1,1))
  gfactors = _apply_winograd_matrix(winograd_G, g, len(HW)).reshape(*HWI, 1, groups, rcout, cin, *([1]*len(tyx)))
  # (HWI, bs, cin_, tyx) -> (HWI, bs, groups, 1 ,cin, *tyx)
  dfactors = _apply_winograd_matrix(winograd_Bt, d, len(HW)).reshape(*HWI, bs, groups, 1, cin, *tyx)

  # matmul; sum across cin: (HWI, bs, groups, rcout, *tyx); then HWI -> HWO: (HWO, bs, groups, rcout, *tyx)
  ret = _apply_winograd_matrix(winograd_At, (gfactors * dfactors).sum(axis=-1-len(HW), dtype=dtype), len(HW))

  # interleave tyx and HWO: (bs, groups, rcout, oy, HO, ox, WO)
  ret = ret.permute([*range(len(HW), len(ret.shape)-len(HW)), *[i+o for i in range(len(HW)) for o in [len(ret.shape)-len(HW),0]]])
  # merge groups and rcout, tyx and HWO: (bs, groups, cout, *yx), shrink to final
  ret = ret.reshape(bs, cout, *[c * HWO[i] for i, c in enumerate(tyx)]).shrink(tuple((0, s) for s in [bs, cout, *oyx]))

  return (ret if bias is None else ret.add(bias.reshape(1, -1, *[1 for _ in range(len(HW))]))).contiguous().contiguous_backward()