[Bug] 'tvm.relax.op.nn' has no attribute 'attention_bias'

[Cookiee235]

Actual behavior

@I.ir_module
class Module:
    @R.function
    def main(q: R.Tensor((4, 16, 32, 8), dtype="float32"), k: R.Tensor((4, 8, 32, 8), dtype="float32"), v: R.Tensor((4, 8, 32, 16), dtype="float32"), bias: R.Tensor((4, 32, 16, 8), dtype="float32")) -> R.Tensor((4, 16, 32, 16), dtype="float32"):
        gv: R.Tensor((4, 16, 32, 16), dtype="float32") = R.nn.attention_bias(q, k, v, bias, scale=T.float32(0.10000000000000001), causal_mask="TopLeft", window_size=None)
        return gv

error: module 'tvm.relax.op.nn' has no attribute 'attention_bias'
 --> <str>:9:58
   |  
 9 |          gv: R.Tensor((4, 16, 32, 16), dtype="float32") = R.nn.attention_bias(q, k, v, bias, scale=T.float32(0.10000000000000001), causal_mask="TopLeft", window_size=None)
   |                                                           ^^^^^^^^^^^^^^^^^^^ 

Steps to reproduce

irs= """# from tvm.script import ir as I
# from tvm.script import tir as T
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function
    def main(q: R.Tensor((4, 16, 32, 8), dtype="float32"), k: R.Tensor((4, 8, 32, 8), dtype="float32"), v: R.Tensor((4, 8, 32, 16), dtype="float32"), bias: R.Tensor((4, 32, 16, 8), dtype="float32")) -> R.Tensor((4, 16, 32, 16), dtype="float32"):
        gv: R.Tensor((4, 16, 32, 16), dtype="float32") = R.nn.attention(q, k, v, bias, scale=T.float32(0.10000000000000001), causal_mask="TopLeft", window_size=None)
        return gv

import tvm
mod = tvm.script.from_source(irs)
mod.show()
mod_new = tvm.script.from_source(mod.script())   # crash!
"""

cc @Lunderberg @junrushao @Hzfengsy @tqchen

[Cookiee235]

@parsifal-47 This is a similar bug with https://github.com/apache/tvm/pull/17567. Could you help me fix it? Thanks a lot!

[parsifal-47]

@Cookiee235 yes, will be happy to