MCPcopy Create free account
hub / github.com/apache/tvm / main

Method main

tests/python/relax/test_frontend_from_exported_program.py:1331–1350  ·  view source on GitHub ↗
(input: R.Tensor((10, 10), dtype="float32"))

Source from the content-addressed store, hash-verified

1329 class expected_tril:
1330 @R.function
1331 def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
1332 R.Tensor((10, 10), dtype="float32")
1333 ):
1334 # block 0
1335 with R.dataflow():
1336 lv: R.Tensor((10,), dtype="int64") = R.arange(
1337 R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64"
1338 )
1339 lv1: R.Tensor((1, 10), dtype="int64") = R.expand_dims(lv, axis=[-2])
1340 lv2: R.Tensor((10,), dtype="int64") = R.arange(
1341 R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64"
1342 )
1343 lv3: R.Tensor((10, 1), dtype="int64") = R.expand_dims(lv2, axis=[-1])
1344 lv4: R.Tensor((10, 10), dtype="int64") = R.subtract(lv1, lv3)
1345 lv5: R.Tensor((10, 10), dtype="bool") = R.less_equal(lv4, R.const(1, "int64"))
1346 lv6: R.Tensor((), dtype="float32") = R.const(0.0, "float32")
1347 lv7: R.Tensor((10, 10), dtype="float32") = R.where(lv5, input, lv6)
1348 gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv7,)
1349 R.output(gv)
1350 return gv
1351
1352 verify_model(Tril(), example_args, {}, expected_tril)
1353

Callers

nothing calls this directly

Calls 8

TensorMethod · 0.80
dataflowMethod · 0.80
subtractMethod · 0.80
less_equalMethod · 0.80
whereMethod · 0.80
outputMethod · 0.80
arangeMethod · 0.45
expand_dimsMethod · 0.45

Tested by

no test coverage detected