import math, unittest
from dataclasses import replace
from tinygrad import Tensor, dtypes, Context
from tinygrad.uop.ops import ParamArg, UOp, UPat, Ops, PatternMatcher, graph_rewrite

_strip_unique_pm = PatternMatcher([
  (UPat(Ops.BUFFER, name="b"), lambda b: b.replace(arg=replace(b.arg, slot=0)) if isinstance(b.arg, ParamArg) and b.arg.slot != 0 else None),
])
def _strip_unique(u: UOp) -> UOp: return graph_rewrite(u, _strip_unique_pm)

def _t(*shape):
  return Tensor.arange(math.prod(shape)).reshape(*shape)

def _ti(data):  # int32 index tensor
  return Tensor(data, dtype=dtypes.int32)

# Tensor().func().uop should be the same as UOp.func()
def _check(tc: unittest.TestCase, t: Tensor, fn):
  tc.assertIs(fn(t).uop, fn(t.uop), f"\ntensor.uop = {fn(t).uop}\nuop = {fn(t.uop)}")

class TestTensorUOpBinop(unittest.TestCase):
  # Tensor's binop upcasts mixed dtypes via least_upper_dtype + explicit CAST; UOp should match.
  def test_mul_float_int(self):
    t = _t(3).float()
    self.assertIs((t * Tensor.arange(3)).uop, t.uop * UOp.arange(3))
  def test_mul_bool_int(self):
    t = _t(3)
    self.assertIs((t.eq(1) * Tensor.arange(3)).uop, t.uop.eq(1) * UOp.arange(3))
  # Tensor's ufix picks float dtype when scalar is float and self is int; UOp should match.
  def test_add_scalar_float_on_int(self): _check(self, _t(3), lambda x: x + 1.5)
  # div: Tensor.div (default case) delegates to ElementwiseMixin.div; trees must match for Tensor and UOp.
  def test_div_tensor_by_tensor(self):
    a, b = _t(4).float(), _t(4).float() + 1
    self.assertIs((a/b).uop, a.uop/b.uop)
  def test_div_int_by_int(self):                 _check(self, _t(4), lambda x: x / 3)
  def test_div_sum_by_sum(self):                 _check(self, _t(4).float(), lambda x: x.sum() / (x + 1).sum())
  def test_div_broadcast_tensor_by_tensor(self):
    a, b = _t(3, 4).float(), _t(4).float() + 1
    self.assertIs((a/b).uop, a.uop/b.uop)
  # isclose used `self == other` which is Python identity on UOp (not elementwise); now uses .eq().
  def test_isclose(self):
    t = _t(4).float()
    self.assertIs(t.isclose(t).uop, t.uop.isclose(t.uop))
  # __floordiv__/mod/fmod and div(rounding_mode=...) dispatch on dtype in mixin
  def test_floordiv_int(self):   _check(self, _t(4), lambda x: x // 3)
  def test_floordiv_float(self): _check(self, _t(4).float() + 1.5, lambda x: x // 2.0)
  def test_rfloordiv_int(self):  _check(self, _t(4)+1, lambda x: 7 // x)
  def test_mod_int(self):        _check(self, _t(4), lambda x: x % 3)
  def test_mod_float(self):      _check(self, _t(4).float() + 1.5, lambda x: x % 2.0)
  def test_div_trunc_int(self):  _check(self, _t(4), lambda x: x.div(3, rounding_mode="trunc"))
  def test_div_trunc_float(self):_check(self, _t(4).float() + 1.5, lambda x: x.div(2.0, rounding_mode="trunc"))
  def test_fmod_int(self):       _check(self, _t(4), lambda x: x.fmod(3))
  def test_fmod_float(self):     _check(self, _t(4).float() + 1.5, lambda x: x.fmod(2.0))
  def test_floordiv_bool(self):  _check(self, _t(4).cast(dtypes.bool), lambda x: x // True)
  def test_mod_bool(self):       _check(self, _t(4).cast(dtypes.bool), lambda x: x % True)
  def test_fmod_bool(self):      _check(self, _t(4).cast(dtypes.bool), lambda x: x.fmod(True))

class TestTensorUOpClone(unittest.TestCase):
  def test_clone(self):
    t = _t(3, 4).float()
    self.assertIs(_strip_unique(t.clone().uop), _strip_unique(t.uop.clone()))
  def test_clone_deviceless_const(self):
    u = UOp.const(2.0)
    self.assertIs(_strip_unique(Tensor(u).clone().uop), _strip_unique(u.clone()))

class TestTensorUOpGradient(unittest.TestCase):
  def test_gradient(self):
    x = _t(3, 3).float()
    z = (x * 2).sum()
    (tg,) = z.gradient(x)
    (ug,) = z.uop.gradient(x.uop)
    self.assertIs(tg.uop, ug)

class TestTensorUOpGetitem(unittest.TestCase):
  # ---- pure slice patterns ----
  def test_slice_full(self):           _check(self, _t(4), lambda x: x[slice(None)])
  def test_slice_positive(self):       _check(self, _t(8), lambda x: x[1:5])
  def test_slice_open_start(self):     _check(self, _t(8), lambda x: x[:5])
  def test_slice_open_stop(self):      _check(self, _t(8), lambda x: x[3:])
  def test_slice_negative_start(self): _check(self, _t(8), lambda x: x[-3:])
  def test_slice_negative_stop(self):  _check(self, _t(8), lambda x: x[:-2])
  def test_slice_both_negative(self):  _check(self, _t(8), lambda x: x[-5:-1])

  # ---- slice with stride ----
  def test_slice_stride(self):                  _check(self, _t(6), lambda x: x[::2])
  def test_slice_start_stop_stride(self):       _check(self, _t(6), lambda x: x[1:5:2])
  def test_slice_reverse(self):                 _check(self, _t(6), lambda x: x[::-1])
  def test_slice_singleton_negative_step(self): _check(self, _t(8), lambda x: x[3:2:-1])

  # ---- empty / out-of-bounds slice ----
  def test_slice_empty(self):    _check(self, _t(6), lambda x: x[3:1])
  def test_slice_oob_stop(self): _check(self, _t(6), lambda x: x[0:100])

  # ---- single int (reduces a dim) ----
  def test_int_positive(self): _check(self, _t(8), lambda x: x[3])
  def test_int_negative(self): _check(self, _t(8), lambda x: x[-1])

  # ---- ellipsis ----
  def test_ellipsis_only(self):       _check(self, _t(2, 3, 4), lambda x: x[...])
  def test_ellipsis_then_int(self):   _check(self, _t(2, 3, 4), lambda x: x[..., -1])
  def test_ellipsis_then_slice(self): _check(self, _t(2, 3, 4), lambda x: x[..., 1:3])
  def test_ellipsis_then_none(self):  _check(self, _t(2, 3), lambda x: x[..., None])

  # ---- None (unsqueeze) ----
  def test_none_front(self):    _check(self, _t(4), lambda x: x[None])
  def test_none_back(self):     _check(self, _t(4), lambda x: x[:, None])
  def test_none_middle(self):   _check(self, _t(2, 3), lambda x: x[:, None, :])
  def test_multiple_none(self): _check(self, _t(2, 3), lambda x: x[None, :, None])

  # ---- mixed multi-dim ----
  def test_int_then_slice(self):    _check(self, _t(2, 3), lambda x: x[1, :])
  def test_multi_int(self):         _check(self, _t(2, 3, 4), lambda x: x[1, 2])
  def test_mixed_slice_int(self):   _check(self, _t(2, 3, 4), lambda x: x[0:2, -1, 1:3])
  def test_mixed_slice_slice(self): _check(self, _t(3, 4, 5), lambda x: x[1:3, :, 0:2])
  def test_high_rank_combo(self):   _check(self, _t(4, 5, 6), lambda x: x[1:3, :, -1, None])

  # ---- advanced indexing: UOp index on UOp must match Tensor index on Tensor (same uop) ----
  def _check_adv(self, t, idx):
    # idx is a Tensor or a tuple mixing Tensor index arrays with slice/int/None/Ellipsis
    ui = idx.uop if isinstance(idx, Tensor) else (tuple(i.uop if isinstance(i,Tensor) else i for i in idx) if isinstance(idx, tuple) else idx)
    self.assertIs(t[idx].uop, t.uop[ui])

  def test_adv_single(self):           self._check_adv(_t(5), _ti([2,1,0,1,2]))
  def test_adv_negative(self):         self._check_adv(_t(5), _ti([-1,-2,0]))
  def test_adv_out_of_bounds(self):    self._check_adv(_t(5), _ti([4,7,2]))           # oob -> 0
  def test_adv_2d_dim0(self):          self._check_adv(_t(3,4), _ti([2,0,1]))
  def test_adv_2d_index_array(self):   self._check_adv(_t(4,5), _ti([[0,1],[2,3]]))   # shaped index
  def test_adv_two_consecutive(self):  self._check_adv(_t(3,4), (_ti([2,0,1]), _ti([1,2,3])))  # linear path
  def test_adv_two_broadcast(self):    self._check_adv(_t(3,4), (_ti([2,0,1]), _ti([[1],[2],[3]])))
  def test_adv_three_consecutive(self):self._check_adv(_t(3,4,5), (_ti([2,0]), _ti([1,3]), _ti([4,2])))
  def test_adv_after_slice(self):      self._check_adv(_t(2,3,4), (slice(None), _ti([2,0,1])))
  def test_adv_non_consec_permute(self): self._check_adv(_t(2,3,4), (_ti([1,0]), slice(None), _ti([2,0])))
  def test_adv_idx_then_int(self):     self._check_adv(_t(3,4), (_ti([2,0,1]), 2))
  def test_adv_int_then_idx(self):     self._check_adv(_t(3,4), (1, _ti([2,0,1])))
  def test_adv_idx_then_none(self):    self._check_adv(_t(3,4), (_ti([2,0,1]), None))
  def test_adv_none_then_idx(self):    self._check_adv(_t(3,4), (None, _ti([2,0,1])))
  def test_adv_ellipsis_then_idx(self):self._check_adv(_t(2,3,4), (Ellipsis, _ti([2,0,1])))
  def test_adv_idx_slice_mix(self):    self._check_adv(_t(4,5,6), (_ti([1,3]), slice(1,4), _ti([2,0])))

  # bool index is unsupported
  def test_adv_bool_index_rejected(self):
    with self.assertRaises(IndexError): _t(5)[_t(5) > 2]
    with self.assertRaises(IndexError): _t(5).uop[(_t(5) > 2).uop]

  # python list/tuple indices
  def test_adv_python_list(self):
    self.assertIs(_strip_unique(_t(5)[[2,1,0]].uop), _strip_unique(_t(5).uop[[2,1,0]]))
  def test_adv_python_list_negative(self):
    self.assertIs(_strip_unique(_t(5)[[-1,-2,0]].uop), _strip_unique(_t(5).uop[[-1,-2,0]]))
  def test_adv_python_list_nested(self):
    self.assertIs(_strip_unique(_t(3,4)[[0,1],[2,0]].uop), _strip_unique(_t(3,4).uop[[0,1],[2,0]]))
  def test_adv_python_list_2d_index(self):
    self.assertIs(_strip_unique(_t(4,5)[[[0,1],[2,3]]].uop), _strip_unique(_t(4,5).uop[[[0,1],[2,3]]]))

class TestTensorUOpCumalu(unittest.TestCase):
  def test_cumsum_1d(self):       _check(self, _t(5), lambda x: x.cumsum())
  def test_cumsum_2d(self):       _check(self, _t(3, 4), lambda x: x.cumsum(1))
  def test_cumsum_non_last(self): _check(self, _t(3, 4), lambda x: x.cumsum(0))
  def test_cumsum_large(self):    _check(self, _t(600), lambda x: x.cumsum())  # exercises _split_cumalu
  def test_cumprod(self):         _check(self, _t(4), lambda x: x.cumprod(0))

class TestTensorUOpCumMinMax(unittest.TestCase):
  def _check_pair(self, t, fn):
    vt, it = fn(t)
    vu, iu = fn(t.uop)
    self.assertIs(vt.uop, vu)
    self.assertIs(it.uop, iu)
  def test_cummax_1d(self):    self._check_pair(_t(5), lambda x: x.cummax(0))
  def test_cummax_2d(self):    self._check_pair(_t(3, 4), lambda x: x.cummax(1))
  def test_cummax_0d(self):    self._check_pair(_t(1).reshape(()), lambda x: x.cummax(0))
  def test_cummin_1d(self):    self._check_pair(_t(5), lambda x: x.cummin(0))
  def test_cummin_2d(self):    self._check_pair(_t(3, 4), lambda x: x.cummin(1))

class TestTensorUOpArgMinMax(unittest.TestCase):
  def _check(self, t, fn): self.assertIs(fn(t).uop, fn(t.uop))
  def test_argmax(self):       self._check(_t(3, 4), lambda x: x.argmax(axis=1))
  def test_argmax_flat(self):  self._check(_t(3, 4), lambda x: x.argmax())
  def test_argmin(self):       self._check(_t(3, 4), lambda x: x.argmin(axis=0))

class TestTensorUOpSequential(unittest.TestCase):
  def test_sequential(self): _check(self, _t(4), lambda x: x.sequential([lambda y: y * 2, lambda y: y + 1]))

class TestTensorUOpOneHot(unittest.TestCase):
  def test_one_hot(self):
    t = _t(5)
    self.assertIs(t.one_hot(5).uop, t.uop.one_hot(5))

class TestTensorUOpSort(unittest.TestCase):
  def _check(self, t, **kw):
    tv, ti = t.sort(**kw)
    uv, ui = t.uop.sort(**kw)
    self.assertIs(tv.uop, uv)
    self.assertIs(ti.uop, ui)
  def test_sort_1d(self):         self._check(Tensor([0.5, 0.1, 0.3]).float())
  def test_sort_descending(self): self._check(Tensor([0.5, 0.1, 0.3]).float(), descending=True)
  def test_sort_2d(self):         self._check(_t(2, 4).float())
  def test_sort_single(self):     self._check(Tensor([1.0]).float())
  def test_argsort(self):
    t = Tensor([0.5, 0.1, 0.3]).float()
    self.assertIs(t.argsort().uop, t.uop.argsort())
  def test_topk(self):
    t = _t(2, 4).float()
    tv, ti = t.topk(2)
    uv, ui = t.uop.topk(2)
    self.assertIs(tv.uop, uv)
    self.assertIs(ti.uop, ui)

class TestTensorUOpAllclose(unittest.TestCase):
  def test_allclose(self):
    a, b = _t(4).float(), _t(4).float()
    self.assertIs(a.allclose(b).uop, a.uop.allclose(b.uop))

class TestTensorUOpCast(unittest.TestCase):
  def test_cast_str_dtype(self):
    t = _t(4)
    self.assertIs(t.cast("float32").uop, t.uop.cast("float32"))
    self.assertIs(t.uop.cast("float32").dtype, dtypes.float32)

class TestTensorUOpBitcast(unittest.TestCase):
  def test_bitcast_same_dtype(self): _check(self, _t(4).float(), lambda x: x.bitcast(dtypes.float32))
  def test_bitcast_str_dtype(self):
    t = _t(4)
    self.assertIs(t.bitcast("uint32").uop, t.uop.bitcast("uint32"))
    self.assertIs(t.uop.bitcast("uint32").dtype, dtypes.uint32)
  def test_bitcast_same_and_diff_size(self):
    _check(self, _t(4).float(), lambda x: x.bitcast(dtypes.uint32))              # same size
    _check(self, _t(4).cast(dtypes.uint8), lambda x: x.bitcast(dtypes.uint16))   # widen: uint8[4] -> uint16[2]
    _check(self, _t(4).cast(dtypes.uint16), lambda x: x.bitcast(dtypes.uint8))   # narrow: uint16[4] -> uint8[8]

class TestTensorUOpRand(unittest.TestCase):
  def test_random_bits(self):
    k = UOp.empty((2,), dtype=dtypes.uint32)
    c = UOp.zeros(2, dtype=dtypes.uint32)
    for num in (1, 4, 7, 1024):
      self.assertIs(Tensor.random_bits(Tensor(k), Tensor(c), num).uop, UOp.random_bits(k, c, num))
  def test_bits_to_rand_float32(self):
    bits_uop = UOp.empty((8,), dtype=dtypes.uint32)
    for shape in ((8,), (2, 4), (5,)):
      self.assertIs(Tensor._bits_to_rand(Tensor(bits_uop), shape, dtypes.float32).uop, UOp._bits_to_rand(bits_uop, shape, dtypes.float32))
  def test_threefry(self):
    t = _t(4).cast(dtypes.uint64)
    self.assertIs(t.threefry(t).uop, t.uop.threefry(t.uop))
  def test_threefry_random_bits(self):
    key, c0, c1 = UOp.empty((2,), dtype=dtypes.uint32), UOp.arange(4, dtype=dtypes.uint32), UOp.arange(4, dtype=dtypes.uint32)
    self.assertIs(Tensor._threefry_random_bits(Tensor(key), Tensor(c0), Tensor(c1)).uop, UOp._threefry_random_bits(key, c0, c1))
  def test_rand(self):
    k, c = UOp.empty((2,), dtype=dtypes.uint32), UOp.zeros(2, dtype=dtypes.uint32)
    self.assertIs(Tensor._rand(Tensor(k), Tensor(c), (2, 2), dtypes.float32).uop, UOp._rand(k, c, (2, 2), dtypes.float32))
    self.assertIs(Tensor._rand(Tensor(k), Tensor(c), (0, 3), dtypes.float32).uop, UOp._rand(k, c, (0, 3), dtypes.float32))

class TestTensorUOpGather(unittest.TestCase):
  def _check(self, t, dim, idx):
    self.assertIs(t.gather(dim, idx).uop, t.uop.gather(dim, idx.uop))
  def test_gather_1d(self):  self._check(_t(5), 0, Tensor([2, 1, 0, 1, 2], dtype=dtypes.int32))
  def test_gather_dim0(self): self._check(_t(3, 4), 0, Tensor([[0, 1, 2, 0], [1, 2, 0, 1], [2, 0, 1, 2]], dtype=dtypes.int32))
  def test_gather_dim1(self): self._check(_t(3, 4), 1, Tensor([[0, 1, 2, 3], [1, 2, 3, 0], [2, 3, 0, 1]], dtype=dtypes.int32))

class TestTensorUOpInterpolate(unittest.TestCase):
  def _check(self, t, mode):
    self.assertIs(t.interpolate(size=(2, 2), mode=mode).uop, t.uop.interpolate(size=(2, 2), mode=mode))
  def test_interpolate_nearest(self):       self._check(_t(1, 1, 4, 4).float(), "nearest")
  def test_interpolate_nearest_exact(self): self._check(_t(1, 1, 4, 4).float(), "nearest-exact")
  def test_interpolate_linear(self):        self._check(_t(1, 1, 4, 4).float(), "linear")

class TestTensorUOpLoss(unittest.TestCase):
  def test_cross_entropy(self):
    t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32)
    self.assertIs(t.cross_entropy(Y).uop, t.uop.cross_entropy(Y.uop))
  def test_sparse_categorical_crossentropy(self):
    t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32)
    self.assertIs(t.sparse_categorical_crossentropy(Y).uop, t.uop.sparse_categorical_crossentropy(Y.uop))
  def test_sparse_categorical_crossentropy_ignore_index(self):
    t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32)
    self.assertIs(t.sparse_categorical_crossentropy(Y, ignore_index=0).uop, t.uop.sparse_categorical_crossentropy(Y.uop, ignore_index=0))
  def test_nll_loss(self):
    t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
    self.assertIs(t.nll_loss(Y).uop, t.uop.nll_loss(Y.uop))
  def test_nll_loss_weight(self):
    t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float()
    self.assertIs(t.nll_loss(Y, weight=w).uop, t.uop.nll_loss(Y.uop, weight=w.uop))
  def test_nll_loss_ignore_index(self):
    t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
    self.assertIs(t.nll_loss(Y, ignore_index=1).uop, t.uop.nll_loss(Y.uop, ignore_index=1))
  def test_nll_loss_none_reduction(self):
    t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
    self.assertIs(t.nll_loss(Y, reduction="none").uop, t.uop.nll_loss(Y.uop, reduction="none"))
  def test_nll_loss_weight_ignore_index(self):
    t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float()
    self.assertIs(t.nll_loss(Y, weight=w, ignore_index=1).uop, t.uop.nll_loss(Y.uop, weight=w.uop, ignore_index=1))

class TestTensorUOpScatter(unittest.TestCase):
  def test_scatter(self):
    x, idx, src = _t(3, 4).float(), Tensor([[0, 1, 2, 0]], dtype=dtypes.int32), _t(1, 4).float()
    self.assertIs(x.scatter(0, idx, src).uop, x.uop.scatter(0, idx.uop, src.uop))
  def test_scatter_scalar_src(self):
    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
    self.assertIs(x.scatter(1, idx, 3.14).uop, x.uop.scatter(1, idx.uop, 3.14))
  # inf cannot be cast to int — this regresses if scalar src is routed through index.dtype first
  def test_scatter_inf_src(self):
    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
    self.assertIs(x.scatter(1, idx, float("inf")).uop, x.uop.scatter(1, idx.uop, float("inf")))
  def test_scatter_add(self):
    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
    self.assertIs(x.scatter(1, idx, 3.14, reduce="add").uop, x.uop.scatter(1, idx.uop, 3.14, reduce="add"))
  def test_scatter_multiply(self):
    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
    self.assertIs(x.scatter(1, idx, 3.14, reduce="multiply").uop, x.uop.scatter(1, idx.uop, 3.14, reduce="multiply"))
  # tensor src with reduce hits the "elif reduce: raise" branch in both Tensor and UOp paths
  def test_scatter_tensor_src_with_reduce_raises(self):
    x, idx, src = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32), _t(1, 2).float()
    with self.assertRaises(TypeError): x.scatter(1, idx, src, reduce="add")
    with self.assertRaises(TypeError): x.uop.scatter(1, idx.uop, src.uop, reduce="add")

class TestTensorUOpScatterReduce(unittest.TestCase):
  def _check(self, x, idx, src, **kw):
    self.assertIs(x.scatter_reduce(0, idx, src, **kw).uop, x.uop.scatter_reduce(0, idx.uop, src.uop, **kw))
  def test_sum(self):  self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="sum")
  def test_prod(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="prod")
  def test_mean(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="mean")
  def test_amax(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="amax")
  def test_amin(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="amin")
  def test_mean_exclude_self(self):
    self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="mean", include_self=False)

class TestTensorUOpMaskedSelect(unittest.TestCase):
  # only the fixed-size path is pure
  def _check(self, t, mask, **kw):
    self.assertIs(t.masked_select(mask, **kw).uop, t.uop.masked_select(mask.uop, **kw))
  def test_masked_select_1d(self): self._check(_t(6), Tensor([True, False, True, False, True, False]), size=4)
  def test_masked_select_2d(self):
    self._check(_t(3, 3), Tensor([[True, False, True], [False, True, False], [False, False, True]]), size=6, fill_value=-1)

class TestTensorUOpNonzero(unittest.TestCase):
  def _check(self, t, **kw): self.assertIs(t.nonzero(**kw).uop, t.uop.nonzero(**kw))
  def test_nonzero_1d(self): self._check(_t(5), size=3)
  def test_nonzero_2d(self): self._check(_t(2, 3), size=4)

class TestTensorUOpPool(unittest.TestCase):
  def test_avg_pool2d(self):                _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d())
  def test_avg_pool2d_padding(self):        _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d(padding=1))
  def test_avg_pool2d_ceil(self):           _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d(ceil_mode=True))
  def test_avg_pool2d_no_count_pad(self):   _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d(padding=1, count_include_pad=False))
  def test_max_pool2d(self):                _check(self, _t(1, 1, 5, 5).float(), lambda x: x.max_pool2d())
  def test_max_pool2d_padding(self):        _check(self, _t(1, 1, 5, 5).float(), lambda x: x.max_pool2d(padding=1))
  def test_max_pool2d_ceil(self):           _check(self, _t(1, 1, 5, 5).float(), lambda x: x.max_pool2d(ceil_mode=True))
  def test_max_pool2d_return_indices(self):
    t = _t(1, 1, 5, 5).float()
    vt, it = t.max_pool2d(return_indices=True)
    vu, iu = t.uop.max_pool2d(return_indices=True)
    self.assertIs(vt.uop, vu)
    self.assertIs(it.uop, iu)
  def test_max_unpool2d(self):
    t = _t(1, 1, 4, 4).float()
    out, idx = t.max_pool2d(return_indices=True)
    self.assertIs(out.max_unpool2d(idx).uop, out.uop.max_unpool2d(idx.uop))

class TestTensorUOpCat(unittest.TestCase):
  def test_cat_dim0(self):     _check(self, _t(2, 3), lambda x: x.cat(x, dim=0))
  def test_cat_dim1(self):     _check(self, _t(2, 3), lambda x: x.cat(x, dim=1))
  def test_cat_3tensors(self): _check(self, _t(2, 3), lambda x: x.cat(x, x, dim=0))
  def test_cat_neg_dim(self):  _check(self, _t(2, 3, 4), lambda x: x.cat(x, dim=-1))

class TestTensorUOpPad(unittest.TestCase):
  def test_pad_flat(self):               _check(self, _t(4, 5), lambda x: x.pad((1, 2, 0, 3)))
  def test_pad_flat_negative(self):      _check(self, _t(4, 5), lambda x: x.pad((1, -1, 0, 2), value=-1.0))
  def test_pad_grouped_none(self):       _check(self, _t(4, 5), lambda x: x.pad((None, (0, 3))))
  def test_pad_circular(self):           _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="circular"))
  def test_pad_circular_zero_after(self):_check(self, _t(4, 5), lambda x: x.pad(((1, 0), (2, 0)), mode="circular"))
  def test_pad_reflect(self):            _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="reflect"))
  def test_pad_reflect_negative(self):   _check(self, _t(4, 5), lambda x: x.pad(((1, -1), (0, 2)), mode="reflect"))
  def test_pad_replicate(self):          _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="replicate"))
  def test_pad_replicate_negative(self): _check(self, _t(4, 5), lambda x: x.pad(((1, -1), (0, 2)), mode="replicate"))

class TestTensorUOpStack(unittest.TestCase):
  def test_stack_dim0(self):     _check(self, _t(2, 3), lambda x: x.stack(x, dim=0))
  def test_stack_dim1(self):     _check(self, _t(2, 3), lambda x: x.stack(x, dim=1))
  def test_stack_3tensors(self): _check(self, _t(2, 3), lambda x: x.stack(x, x, dim=0))
  def test_stack_new_last(self): _check(self, _t(2, 3), lambda x: x.stack(x, dim=-1))
  def test_stack_mixed_dtype(self):
    w = _t(2, 3).float()
    _check(self, _t(2, 3), lambda x: x.stack(w if isinstance(x, Tensor) else w.uop))
    self.assertIs(_t(2, 3).uop.stack(w.uop).dtype, dtypes.float32)
  def test_stack_index_dtype(self):
    # index is outside the promotion lattice, equal dtypes bypass promotion
    self.assertEqual(UOp.const(1).stack(UOp.const(2)).shape, (2,))

class TestTensorUOpConv2d(unittest.TestCase):
  def test_conv2d_basic(self):
    w = _t(1, 1, 2, 2).float()
    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
  def test_conv2d_padded(self):
    w = _t(1, 1, 2, 2).float()
    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, padding=1))
  def test_conv2d_negative_padding(self):
    w = _t(1, 1, 3, 3).float()
    _check(self, _t(1, 1, 5, 5).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, padding=(-1,-1,-1,-1)))
  def test_conv2d_multichannel_bias(self):
    w, b = _t(4, 2, 3, 3).float(), _t(4).float()
    _check(self, _t(2, 2, 5, 5).float(), lambda x: x.conv2d(*(y if isinstance(x, Tensor) else y.uop for y in (w, b))))
  def test_conv2d_stride_dilation(self):
    w = _t(2, 2, 2, 2).float()
    _check(self, _t(1, 2, 6, 6).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, stride=2, dilation=2))
  def test_conv2d_groups(self):
    w = _t(4, 1, 2, 2).float()
    _check(self, _t(1, 4, 4, 4).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, groups=4))
  def test_conv2d_3d(self):
    w = _t(1, 1, 2, 2, 2).float()
    _check(self, _t(1, 1, 3, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
  def test_conv2d_winograd(self):
    w, a = _t(2, 2, 3, 3).float(), _t(1, 2, 6, 6).float()
    with Context(WINO=0): direct = a.conv2d(w).uop
    with Context(WINO=1):
      self.assertIsNot(a.conv2d(w).uop, direct)
      _check(self, a, lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
  def test_conv2d_image(self):
    w, a = _t(4, 4, 3, 3).float(), _t(1, 4, 8, 8).float()
    with Context(IMAGE=0): direct = a.conv2d(w).uop
    with Context(IMAGE=1):
      self.assertIsNot(a.conv2d(w).uop, direct)
      _check(self, a, lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
  def test_dot_image(self):
    y, a = _t(4, 3).float(), _t(2, 4).float()
    with Context(IMAGE=0): direct = a.dot(y).uop
    with Context(IMAGE=1):
      self.assertIsNot(a.dot(y).uop, direct)
      _check(self, a, lambda x: x.dot(y if isinstance(x, Tensor) else y.uop))
  def test_conv_transpose2d_basic(self):
    w = _t(1, 1, 2, 2).float()
    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv_transpose2d(w if isinstance(x, Tensor) else w.uop))
  def test_conv_transpose2d_stride(self):
    w = _t(1, 1, 2, 2).float()
    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv_transpose2d(w if isinstance(x, Tensor) else w.uop, stride=2))

class TestTensorUOpHashing(unittest.TestCase):
  def test_keccak_sha3_256(self):  _check(self, _t(8).cast(dtypes.uint8), lambda x: x.keccak())
  def test_keccak_shake_128(self): _check(self, _t(8).cast(dtypes.uint8), lambda x: x.keccak("shake_128"))

class TestTensorUOpEinsum(unittest.TestCase):
  def test_einsum_dot(self):       _check(self, _t(2, 3), lambda x: type(x).einsum("ij,ij->", x, x))
  def test_einsum_transpose(self): _check(self, _t(2, 3), lambda x: type(x).einsum("ij->ji", x))

class TestTensorUOpSoftmax(unittest.TestCase):
  def test_softmax_default(self):     _check(self, _t(2, 3).float(), lambda x: x.softmax())
  def test_softmax_axis0(self):       _check(self, _t(2, 3).float(), lambda x: x.softmax(axis=0))
  def test_log_softmax_default(self): _check(self, _t(2, 3).float(), lambda x: x.log_softmax())
  def test_log_softmax_axis0(self):   _check(self, _t(2, 3).float(), lambda x: x.log_softmax(axis=0))

class TestTensorUOpQR(unittest.TestCase):
  def _check(self, t):
    qt, rt = t.qr()
    qu, ru = t.uop.qr()
    self.assertIs(qt.uop, qu)
    self.assertIs(rt.uop, ru)
  def test_qr_square(self):   self._check(_t(3, 3).float())
  def test_qr_tall(self):     self._check(_t(4, 3).float())
  def test_qr_wide(self):     self._check(_t(3, 4).float())
  def test_qr_zero_col(self): self._check(Tensor([[0.0, 1.0], [0.0, 2.0]]))
  def test_qr_batched(self):  self._check(_t(2, 3, 3).float())

class TestTensorUOpSVD(unittest.TestCase):
  def _check(self, t, **kw):
    ut, st, vt = t.svd(**kw)
    uu, su, vu = t.uop.svd(**kw)
    self.assertIs(ut.uop, uu)
    self.assertIs(st.uop, su)
    self.assertIs(vt.uop, vu)
  def test_svd_square(self):    self._check(_t(2, 2).float())
  def test_svd_tall(self):      self._check(_t(3, 2).float())
  def test_svd_wide(self):      self._check(_t(2, 3).float())
  def test_svd_odd_num(self):   self._check(_t(3, 3).float())  # exercises odd-num runoff path
  def test_svd_batched(self):   self._check(_t(2, 2, 2).float())
  def test_svd_nonfull(self):   self._check(_t(3, 2).float(), full_matrices=False)

class TestUOpEmpty(unittest.TestCase):
  def test_empty_dtype_string(self):
    self.assertEqual(UOp.empty((3, 4), dtype="float32").dtype, dtypes.float32)

  def test_empty_like_dtype_override(self):
    u = Tensor.ones(3, 4).uop.empty_like(dtype=dtypes.int8)
    self.assertEqual((u.shape, u.dtype), ((3, 4), dtypes.int8))
    self.assertTrue(u.has_buffer_identity())

  def test_empty_like_sharded_to_single_device(self):
    # regression: sharded source, override to single device must yield full logical shape with no axis
    t = Tensor.ones(8, 4).shard(("NULL:0", "NULL:1"), axis=0)
    for dev in ("NULL:2", ("NULL:2",)):  # singleton tuple also canonicalizes to single device
      u = t.uop.empty_like(device=dev, dtype=dtypes.int32)
      self.assertEqual((u.shape, u.device, u.dtype, u.axis), ((8, 4), "NULL:2", dtypes.int32, None))
      self.assertTrue(u.has_buffer_identity())

  def test_empty_direct_singleton_tuple_device(self):
    u = UOp.empty((4,), dtype=dtypes.float32, device=("NULL:0",))
    self.assertEqual((u.shape, u.device, u.axis), ((4,), "NULL", None))

class TestTensorUOpCreation(unittest.TestCase):
  def test_empty(self):
    self.assertIs(_strip_unique(Tensor.empty(2, 3).uop), _strip_unique(UOp.empty(2, 3)))
  def test_full(self):
    self.assertIs(_strip_unique(Tensor.full((2, 3), 42).uop), _strip_unique(UOp.full((2, 3), 42)))
  def test_full_kwargs(self):
    self.assertIs(_strip_unique(Tensor.full((2, 3), 42, dtype=dtypes.int8, device="NULL").uop),
                  _strip_unique(UOp.full((2, 3), 42, dtype=dtypes.int8, device="NULL")))
  def test_full_symbolic_fill(self):
    t = Tensor.full((2, 3), UOp.variable("x", 1, 10).bind(5))
    self.assertEqual(t.shape, (2, 3))
  def test_zeros(self):
    self.assertIs(_strip_unique(Tensor.zeros(2, 3).uop), _strip_unique(UOp.zeros(2, 3)))
  def test_ones(self):
    self.assertIs(_strip_unique(Tensor.ones(2, 3).uop), _strip_unique(UOp.ones(2, 3)))
  def test_invalids(self):
    self.assertIs(_strip_unique(Tensor.invalids(2, 3, dtype=dtypes.int8).uop), _strip_unique(UOp.invalids((2, 3), dtype=dtypes.int8)))
  def test_empty_like(self):
    t = Tensor.empty(2, 3, dtype=dtypes.int8)
    self.assertIs(_strip_unique(t.empty_like().uop), _strip_unique(t.uop.empty_like()))
    self.assertIs(_strip_unique(t.empty_like(dtype=dtypes.float, device="NULL").uop), _strip_unique(t.uop.empty_like(dtypes.float, "NULL")))
  def test_arange(self):
    self.assertIs(Tensor.arange(5).uop, UOp.arange(5))
  def test_arange_empty(self):
    self.assertIs(Tensor.arange(5, 5).uop, UOp.arange(5, 5))
  def test_arange_step(self):
    self.assertIs(Tensor.arange(5, 10, 2).uop, UOp.arange(5, 10, 2))
  def test_linspace(self):
    self.assertIs(Tensor.linspace(0, 10, 5).uop, UOp.linspace(0, 10, 5))
  def test_linspace_one_step(self):
    self.assertIs(Tensor.linspace(5, 10, 1).uop, UOp.linspace(5, 10, 1))
  def test_eye(self):
    self.assertIs(Tensor.eye(3).uop, UOp.eye(3))
  def test_eye_rect(self):
    self.assertIs(Tensor.eye(2, 4).uop, UOp.eye(2, 4))
  def test_triu(self):
    t = _t(3, 4)
    self.assertIs(t.triu().uop, t.uop.triu())
  def test_triu_diagonal(self):
    t = _t(3, 4)
    self.assertIs(t.triu(diagonal=1).uop, t.uop.triu(diagonal=1))
  def test_tril(self):
    t = _t(3, 4)
    self.assertIs(t.tril().uop, t.uop.tril())
  def test_tril_diagonal(self):
    t = _t(3, 4)
    self.assertIs(t.tril(diagonal=-1).uop, t.uop.tril(diagonal=-1))

if __name__ == "__main__":
  unittest.main()
