sonnet abstract module in tensorflow2

Hello friends

I want to use i3d model created by sonnet in tensorflow 2

these are my environment:
ubuntu 20
tensorflow 2.12
python 3.8

this is my i3d code:

`
from __future__ import absolute_import

from __future__ import division
from __future__ import print_function
import sonnet as snt
import tensorflow as tf
class Unit3D(snt.AbstractModule):

`

but when I want to import it in tf2, I face the following error:
AttributeError: module 'sonnet' has no attribute 'AbstractModule'

which is related to this line:

class Unit3D(snt.AbstractModule):
>>> class MultiplyModule(snt.Module):
... def __call__(self, x):
... if not hasattr(self, 'w'):
... self.w = tf.Variable(2., name='w')
... return x * self.w

>>> mod = MultiplyModule()
>>> mod(1.)

how to fix it ?