torchrec
torchrec copied to clipboard
pytorch
register custom_op for fpEBC
Summary:
context
- convert
FeatureProcessedEmbeddingBagCollectionto custom op in IR export - add serialization and deserialization function for FPEBC
- add an API for the
FeatureProcessorInterfaceto export necessary paramters for create an instance - use this API (
get_init_kwargs) in the serialize and deserialize functions to flatten and unflatten the feature processor
details
- Added
FPEBCMetadataschema for FP_EBC, use afp_jsonstring to store the necessary paramters - Added
FPEBCJsonSerializer, converted the init_kwargs to json string and store in thefp_jsonfield in the metadata - Added a fqn check for
serialized_fqns, so that when a higher-level module is serialized, the lower-level module can be skipped (it's already included in the higher-level module) - Added an API called
get_init_kwargsforFeatureProcessorsCollectionandFeatureProcessor, and use aFeatureProcessorNameMapto map the classname to the feature processor class - Added
_non_strict_exporting_forwardfunction for FPEBC so that in non_strict IR export it goes to the custom_op logic
Differential Revision: D57829276
