Skip to content

Commit f6f9304

Browse files
Alfie-Edwardsgeorgepaw
authored andcommitted
Adding ModelExtensions to ipu.keras.extensions namespace
Summary: Adding ModelExtensions to the ipu.keras.extensions namespace. SequentialExtensions and FuntionalExtensions are both in there. TF2.5 Only Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, georgep Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, georgep Maniphest Tasks: T58378 Differential Revision: https://phabricator.sourcevertex.net/D63220
1 parent 568a1ea commit f6f9304

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

tensorflow/python/ipu/keras/extensions/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,4 +23,6 @@
2323
from tensorflow.python.ipu.keras.extensions.functional_extensions import FunctionalExtension
2424
from tensorflow.python.ipu.keras.extensions.sequential_extensions import SequentialLayerPipelineStageAssignment
2525
from tensorflow.python.ipu.keras.extensions.sequential_extensions import SequentialExtension
26+
from tensorflow.python.ipu.keras.extensions.model_extensions import ModelLayerPipelineStageAssignment
27+
from tensorflow.python.ipu.keras.extensions.model_extensions import ModelExtension
2628
# pylint: enable=unused-import

0 commit comments

Comments
 (0)