Method test_softmax_backward

test/null/test_schedule.py:808–812

(self)

    self.assertEqual(linear.src[2].src[1].dtype, dtypes.float)

  def test_softmax_backward(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 12, 64, 64).realize()
    x.softmax().sum().backward()
    run_linear(*check_schedule(x.grad, 4))

  def test_scaled_dot_product_attention_fusion(self):
    x, y, z, m = (Tensor.empty(32, 8, 16, 16) for _ in range(4))