tinygrad/extra/optimization/test_beam_search.py
chenyu 0ca0e9ee5e
exclude ast with variables from beam search (#2140)
* exclude ast with variables from beam search

* test that

* add to CI
2023-10-25 16:35:29 -04:00

16 lines
428 B
Python

import unittest
from tinygrad.helpers import BEAM
from tinygrad.shape.symbolic import Variable
from tinygrad.tensor import Tensor
class TestBeamSearch(unittest.TestCase):
def setUp(self):
self.old_beam = BEAM.value
BEAM.value = 2
def tearDown(self):
BEAM.value = self.old_beam
def test_variable_ast_no_beam(self):
a = Tensor.rand(3, 3).reshape((Variable("a", 1, 10).bind(3), 3))
a = (a+1).realize()