[JAX] Adjust Module Structure.

[mingxu1067]

1. Collect Flax related modules to a sub-folder, flax.
2. Add a function to unify scale_init for zero-centered-gamma LN.

Note: This PR is for the further Praxis (another NN library on top of JAX) support.

[mingxu1067]

/te-ci

[nouiz]

LGTM

[nouiz]

Should we update the JAX examples?

[timmoon10]

Seems reasonable. It's a big API change, so it's best to do it before this has many downstream users. We need to update the documentation as well as the examples to use this new API, e.g. sed 's/te.DenseGeneral/te.flax.DenseGeneral/'.

[mingxu1067]

Update:

1. Adding flax modules to transformer_engine.jax.__init__.py to make this downward compatible.
2. Adapt this changes to jax/examples.
3. Update documents and add deprecated warning.

@ptrendx and @timmoon10 , could we have this changes in v0.8.0 with deprecated warning, then offically move all modules into flax sub-folder in v0.9.0

[mingxu1067]

/te-ci

[ptrendx]

What about the example in README?

[mingxu1067]

Update:

1. Update README.rst.
2. Add deprecate_wrapper in transformer_engine.common.utils.py for general deprecated warning.

[mingxu1067]

/te-ci

[timmoon10]

This doesn't feel quite right to me. Usually when interacting with Python enums we are accessing class members, but this changes it so we interact with the members of a specific instance of DeprecatedEnum. It seems we are kind of reimplementing metaclasses here, and doing it "properly" would require delving into the depths of advanced Python features...

Not a blocker BTW. We don't need to be 100% robust since this is a backward compatibility convenience, and I suspect doing it "right" would take more effort than it's worth.

[mingxu1067]

I agree this is kind of a workaround since I have not found any good approaches to set MetaClass after classes defination.
What we need here is

1. Give deprecated warning when using Enum via deprecated way (transformer_engine.jax.TransformerLayerType in this case).
2. No impact to usage.

Reimplementing metaclasses can break the requirement 1, for example:

classDeprecatedEnumMeta(type):
def__new__(cls, name, bases, dct):
warnings.warn(self.msg, DeprecationWarning)
# Adding TransformerLayerEnum's members into dctclassDeprecateEnum(metaclass=DeprecatedEnumMeta):
pass

The warning message would be showed right after DeprecateEnum defined, instead of invoking (DeprecateEnum.xxxx), which is equivalent to give warning in importing stage.

To my knowledge, there is no way to catch the access of class's members (something like __getattr__ but for classes not instances). Then I implent this as a workaround.

For no impact to usage, it would be great to implement __iter__(), like

classDeprecatedEnum: # pylint: disable=too-few-public-methods"""DeprecatedEnum"""def__init__(self, enum_cls, msg):
self.enum_cls=enum_clsself.msg=msgdef__iter__(self):
returniter(list(self.enum_cls.__members__.values()))
def__getattr__(self, name):
ifnameinself.enum_cls.__members__:
warnings.warn(self.msg, DeprecationWarning)
returnself.enum_cls.__members__[name]
raiseAttributeError(f"{self.enum_cls} does not contain {name}")
>>>importtransformer_engine.jaxaste>>>te.TransformerLayerType.DECODERinte.TransformerLayerTypetransformer_engine/common/utils.py:22: DeprecationWarning: TransformerLayerTypeismovingtotransformer_engine.jax.flaxmodulewarnings.warn(self.msg, DeprecationWarning)
True>>>te.TransformerLayerType.DECODERiste.flax.TransformerLayerType.DECODERTrue

Refer to 5c1f169

[mingxu1067]

Does this make sense?

[timmoon10]

The handling for DeprecatedEnum is starting to get messy. I'm not very familiar with this level of esoteric Python, but I wonder if the following would work:

defmake_deprecated_wrapper(obj, msg):
"""Make a wrapper for a deprecated class or function"""def_warn():
warnings.warn(msg, DeprecationWarning)
ifinspect.isclass(obj):
classDeprecatedMetaclass(type):
"""Metaclass for deprecated class wrapper"""def__getattribute__(cls, name):
_warn()
returnsuper().__getattribute__(cls, name)
classDeprecatedClass(obj, metaclass=DeprecatedMetaclass):
"""Deprecated class wrapper"""def__init__(self, *args, **kwargs):
_warn()
super().__init__(*args, **kwargs)
returnDeprecatedClassifinspect.isfunction(obj):
defdeprecated_function(*args, **kwargs):
_warn()
returnobj(*args, **kwargs)
returndeprecated_functionraiseNotImplementedError(
f"make_deprecated_wrapper supports classes and functions, but got {type(obj)}.")

Thoughts:

• No need to treat enums different than other classes.
• Minimal function overriding and no extra state.
• Printing the warning in a helper function prevents redundant messages if we both create a class instance and access a class variable.
• I am not sure how this will handle the case where you have multiple deprecated classes. I think warnings will filter based on line number, so it will probably only complain about the first deprecated thing you use.

This is mostly for personal interest and isn't a serious suggestion.

[mingxu1067]

The potential issue of above suggestion is that Enum in Python is not allowed to extend. If we define a class to inherit a Enum, then we get cannot extend enumeration xxxx. Let's discuss offline to figure out a better approach for improvment.

[timmoon10]

Interesting, that does complicate things. I guess dealing with enums is going to be messy no matter what then.

[nouiz]

Note, if it is too complicated for enum, we could just not give deprecation warning in the code and only in the doc.
If they use the old enum, they will also use the old fct and classes and those already give a warning. So they will get warning in all cases.
I wouldn't speed too much time on that now.

[ptrendx]

LGTM from my side.

[ksivaman]

/te-ci