[XLA:CPU] Add MpiCollectives to the the .pyi stubs #11721
Closed
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
inailuig
changed the title
hawkinsp
approved these changes
Apr 22, 2024
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Apr 22, 2024
Imported from GitHub PR openxla/xla#11721 we forgot this in openxla/xla#7849. Copybara import of the project: -- 3924cc0fbbb63e9503f38a59aede3b8e817b17fa by Clemens Giuliani <clemens@inailuig.it>: [XLA:CPU] add missing type annotations for the mpi collectives Merging this change closes #11721 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11721 from inailuig:mpicollectives_pytype 3924cc0fbbb63e9503f38a59aede3b8e817b17fa PiperOrigin-RevId: 627055918
sgerrard
approved these changes
Apr 23, 2024
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Apr 23, 2024
Imported from GitHub PR openxla/xla#11721 we forgot this in openxla/xla#7849. Copybara import of the project: -- 3924cc0fbbb63e9503f38a59aede3b8e817b17fa by Clemens Giuliani <clemens@inailuig.it>: [XLA:CPU] add missing type annotations for the mpi collectives Merging this change closes #11721 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11721 from inailuig:mpicollectives_pytype 3924cc0fbbb63e9503f38a59aede3b8e817b17fa PiperOrigin-RevId: 627055918
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Apr 24, 2024
…apped COO format. Reverts 2ba594d FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11721 from inailuig:mpicollectives_pytype 3924cc0fbbb63e9503f38a59aede3b8e817b17fa PiperOrigin-RevId: 623056684
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Apr 24, 2024
Here is the signature of the provide API:
```c++
// Converts a TensorFlow model (either from a SavedModel or an MLIR module) to a
// StableHLO MLIR module.
//
// Args:
// input_path: The path to the input TensorFlow SavedModel or MLIR module.
// context: The MLIR context to use for parsing or creating the MLIR module.
// exported_model_signatures: A comma-separated list of exported model
// signatures (functions) to convert.
// tag_names: A comma-separated list of tag names used for loading SavedModel.
// input_arg_shapes_str: A string representation of input argument shapes.
// Shapes for different tensors are separated by ':', and dimension sizes for
// the same tensor are separated by ','. For example,
// 'input-arg-shapes=1,2::1,?' expresses input arguments with shapes [1,2],
// [] and [1,?].
// is_input_mlir_module: If true, `input_path` is treated as an MLIR
// module instead of a SavedModel.
//
// Returns:
// An absl::StatusOr containing the converted StableHLO MLIR module on
// success, or an absl::Status with an error message on failure.
absl::StatusOr<OwningOpRef<ModuleOp>> TfToStablehlo(
absl::string_view input_path, MLIRContext* context,
absl::string_view exported_model_signatures, absl::string_view tag_names,
absl::string_view input_arg_shapes_str, bool is_input_mlir_module = false);
```
Reverts 2ba594d
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11721 from inailuig:mpicollectives_pytype 3924cc0fbbb63e9503f38a59aede3b8e817b17fa
PiperOrigin-RevId: 625351420
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Apr 24, 2024
Imported from GitHub PR openxla/xla#11721 we forgot this in openxla/xla#7849. Copybara import of the project: -- 3924cc0fbbb63e9503f38a59aede3b8e817b17fa by Clemens Giuliani <clemens@inailuig.it>: [XLA:CPU] add missing type annotations for the mpi collectives Merging this change closes #11721 PiperOrigin-RevId: 627544891
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Apr 24, 2024
1. Added a configuration option to run executable with profiler and collect FDO. See config.py and pxla.py. 2. Added PGLESessionRunner class which is able to run profile session and aggregate results. Each py_executable now have profile_session a property which might be setup by JAX before running a model. 3. Patched JAX jit login to recompile module until FDO would be collected. See pjit.py. 4. Add optional flag to share FDO profile between hosts. Users might choose either share fdo profiles or share an executable after compilation. FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11721 from inailuig:mpicollectives_pytype 3924cc0fbbb63e9503f38a59aede3b8e817b17fa PiperOrigin-RevId: 617824989
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
we forgot this in #7849.