flex_model.distributed.BaseRoutingStrategy

class flex_model.distributed.BaseRoutingStrategy(prologue_fn, epilogue_fn)

Defines a routing strategy, which every device participates in. Moves corresponding tensors via collective communication.

Methods

`__init__`(prologue_fn, epilogue_fn)
`execute_epilogue`(tensor)
`execute_prologue`(tensor)
`initialize`(tensor, expected_shape)