| 34 | namespace grappler { |
| 35 | |
| 36 | Status RebatchOptimizer::Init( |
| 37 | const tensorflow::RewriterConfig_CustomGraphOptimizer* config) { |
| 38 | if (!config) |
| 39 | return errors::InvalidArgument( |
| 40 | "Cannot initialize RebatchOptimizer without config."); |
| 41 | |
| 42 | num_replicas_ = config->parameter_map().at("num_replicas").i(); |
| 43 | use_fallback_ = config->parameter_map().at("use_fallback").b(); |
| 44 | return Status::OK(); |
| 45 | } |
| 46 | |
| 47 | namespace { |
| 48 |