mlx - numpy ops and merge
#21279
Merged
mlx - numpy ops and merge
#21279
Conversation
This file contains hidden or 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
* Add random_posterization processing layer * Add test cases * correct failed case
* Add random_sharpness.py * Update random_sharpness * Add test cases * Fix failed test case
* Add random_shear processing layer * Update method name * Fix failed test case * Fix failed test case * Fix failed test case
Bumps the github-actions group with 2 updates: [actions/upload-artifact](https://github.com/actions/upload-artifact) and [github/codeql-action](https://github.com/github/codeql-action). Updates `actions/upload-artifact` from 4.4.3 to 4.5.0 - [Release notes](https://github.com/actions/upload-artifact/releases) - [Commits](actions/upload-artifact@b4b15b8...6f51ac0) Updates `github/codeql-action` from 3.27.5 to 3.28.0 - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](github/codeql-action@f09c1c0...48ab28a) --- updated-dependencies: - dependency-name: actions/upload-artifact dependency-type: direct:production update-type: version-update:semver-minor dependency-group: github-actions - dependency-name: github/codeql-action dependency-type: direct:production update-type: version-update:semver-minor dependency-group: github-actions ... Signed-off-by: dependabot[bot] <[email protected]> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
* Update BackupAndRestore class example * Update backup_and_restore.py --------- Co-authored-by: François Chollet <[email protected]>
…eam#20710) * Refactor export_lib and add export_onnx Add tf2onnx requirements * Add onnxruntime dep * Update numpy dep * Resolve comments
…am#20725) * Patch tf2onnx to support numpy 2 * Fix warnings * Update export_onnx
….py (keras-team#20728) The `window_length` parameter in `tf.keras.ops.istft` requires `tf.int32` dtype, but this isn't documented. This can cause unexpected `ValueError` when using `tf.int64` and `tf.int16` Here is the Example case: ``` import tensorflow as tf input_dict = { 'stfts': tf.constant([[-0.87817144+1.14583987j, -0.32066484+0.25565411j]], dtype=tf.complex128), 'frame_length': tf.constant(256, dtype=tf.int16), 'frame_step': tf.constant(5120,dtype=tf.int64) } result = tf.signal.inverse_stft(**input_dict) print(result) ``` The code throws the following error: ``` ValueError: window_length: Tensor conversion requested dtype int32 for Tensor with dtype int64 ```
* Add rand_augment init * Update rand_augment init * Add rand_augment * Add NotImplementedError * Add some test cases * Fix failed test case * Update rand_augment * Update rand_augment test * Fix random_rotation bug * Add build method to supress warning. * Add implementation for transform_bboxes
* fixing wrong trainer assumption that batch dim is always the first one in the mesh * need functools partial * lint * fix test failure when distribution=None * lint2 * fix for test failure * added data sharding for 3D+ meshes * lint3 * added @Property for batch_dim_name + refactoring * fix typo
keras-team#20732) The `dtype` / `DTypePolicy` is applied to all float variables.
by simply using `ops.shape(x)` instead of `x.shape`.
The same tests exist at: - https://github.com/keras-team/keras/blob/master/keras/src/export/saved_model_test.py#L66 - https://github.com/keras-team/keras/blob/master/keras/src/export/onnx_test.py#L62 The goal is to isolate the use of `onnxruntime` to a single file, `onnx_test.py`.
* Add OpenVINO into README.md Signed-off-by: Kazantsev, Roman <[email protected]> * Update README.md --------- Signed-off-by: Kazantsev, Roman <[email protected]>
…eam#20738) Multiple Example Title has removed in metrics.MeanIoU method
* Fix JAX GPU CI * Makes formatter happy * Makes formatter happy - 2
In particular for functional models.
…ras-team#20745) * feat(ops): Add keras.ops.image.rot90 operation Adds a new operation to rotate tensors by 90 degrees in the specified plane: - Implements rot90 operation in keras.ops.image module - Adds support for multiple rotations (k parameter) and custom axes - Matches numpy.rot90 behavior and API for consistency - Adds comprehensive test coverage including batch images support - Handles input validation for tensor dimensions and axes - Supports symbolic tensor execution The operation follows the same interface as numpy.rot90 and tf.image.rot90: rot90(array, k=1, axes=(0, 1)) * feat: add JAX, NumPy and PyTorch backends for rot90 Add implementations of rot90() for multiple backend frameworks: - JAX backend implementation - NumPy backend implementation - PyTorch backend implementation * Move rot90 from image to numpy ops Move rot90 operation to numpy.py files in backend implementations since it's a numpy op (https://numpy.org/doc/stable/reference/generated/numpy.rot90.html). Now exported as both keras.ops.rot90 and keras.ops.numpy.rot90. * Fix dtype conflict in PyTorch backend's rot90 function Resolved the 'Invalid dtype: object' error by explicitly using to avoid naming conflicts with the custom function. * Replace experimental NumPy rot90 with core TF ops Replace tf.experimental.numpy.rot90 with core TF ops for XLA compatibility. Use convert_to_tensor for input handling.
* Add bartlett for ops * Update excluded_concrete_tests.txt
* Fix stacked RNN with mask in JAX backend * Add unit test for stacked RNN mask * Fix stacked RNN with mask in Numpy backend * Move unit test to stacked_rnn_cells_test
Bumps the github-actions group with 1 update: [github/codeql-action](https://github.com/github/codeql-action). Updates `github/codeql-action` from 3.28.13 to 3.28.16 - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](github/codeql-action@1b549b9...28deaed) --- updated-dependencies: - dependency-name: github/codeql-action dependency-version: 3.28.16 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: github-actions ... Signed-off-by: dependabot[bot] <[email protected]> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
* Add support for numpy.ravel using OpenVINO opset.reshape * Add support for numpy.ravel using OpenVINO opset.reshape * Add support for numpy.ravel using OpenVINO opset.reshape * Add support for numpy.ravel using OpenVINO opset.reshape * Update excluded_concrete_tests.txt
* Fix message formatting Signed-off-by: Emmanuel Ferdman <[email protected]> * Fix message formatting Signed-off-by: Emmanuel Ferdman <[email protected]> --------- Signed-off-by: Emmanuel Ferdman <[email protected]>
When there are dynamic dimensions, like typically the batch size, `tf.broadcast_dynamic_shape` is not always compilable. Replace with an adhoc implementation for dynamic dimensions where we rely on the broadcast itself to fail when the shapes are not broadcastable. Tested with https://github.com/keras-team/keras-rs/blob/main/examples/listwise_ranking.py on GPU as I was not able to distill a simple reproduction of this.
* Add blackman for ops * Add numpy test for blackman * Add excluded_concrete_tests for blackman * Add NotImplementedError for openvino
…ity. (keras-team#21247) Added optimizer variables were orginally interleaved during `build` prior to keras-team#21232, e.g. `{momentum0, velocity0, momentum1, velocity1, ...}`. In keras-team#21232, the order was changed to non-interleaved for some optimizers, e.g. `{momentum0, momentum1, ..., velocity0, velocity1, ...}`. This broke some custom checkpoint serialization compatibility that relied on the order of variables remaining consistent. Here we modify the base function `add_optimizer_variables(...)` to support creating multiple optimizer variables per training variable, and interleaves creation to restore backward compatibility.
* Add view_as_complex and view_as_real methods and tests * Make op jax compatible * Run linter * Remove outdated test * Remove is_complex() call * Skip tests for openvino * fix backend call * update tests * update tests * Use standardize_dtype for comparisons * Move impl to math.py * Remove math namespace
…team#21256) Change keras-team#21239 broke one use case when the axis dimension is dynamic, the type of the indices is not int32, and the op is run in graph mode. Note that the additional unit tests don't actually cover this.
There were several issues with `keras.utils.plot_model()` / `keras.utils.model_to_dot()` when rendering nested models with `expand_nested=True`: - Sequential models were not expanded - Edges going into nested models would always connect to the first box in the subgraph, which is incorrect because: - Functional models can have multiple inputs and the edge needs to point to the correct input - The destination of the edge can be further nested sub-models like Sequential sub-models - Edges going out of nested models would always connect from the last box in the subgraph, which is incorrect because: - Functional models can have multiple outputs and the edge needs to come from the correct layer - The source of the edge can be further nested in sub-models since there is no "output" box This adds to the integration tests. In particular, it augments the way the topology is verified: - it provides detailed message for dangling edges - it can inspected nested subgraphs - it verifies there are no extra edges that the ones expected - it can verify splits in the graph Visual tests: https://colab.research.google.com/gist/hertschuh/fb7dfdbf4fb31139e33e750ab10aaad4/plot_model-testing.ipynb Fixes keras-team#21119
Keep track of a `global_random_seed`, and ensure it is set when initializing `keras.distribution.initialize(...)`. In multi-host processes in JAX, all processes require consistent random number generation. Otherwise, initializers on different hosts would produce inconsistent values, resulting in both compilation and runtime failures.
* Add openvino backend support for numpy.hstack * Handle non-tuple * Handle non-tuple 2 * Handle non-tuple 3 * Add support for numpy.hstack
Previously the iterations count was incremented in `backend_apply_gradients(...)`. However, this meant that if all variables had `overwrite_with_gradient`, or any other special handling within a custom optimizer, the iteration count would never be increased. The iteration isn't updated via "applying gradients" anyways, so it seems to make more sense to increment it directly in `apply`.
In some cases we need to explicitly specify the layout in JAX (e.g. for sparse-core tables). The previous distribution lib was hard-coded to expect `jax.sharding.Sharding`. Generalized to support explicit Layouts. The previous `distribute_variable`/`distribute_tensor` code also used explicit device placement calls. JAX should be able to handle all this automatically. See [Making process-spanning arrays from external data](https://docs.jax.dev/en/latest/multi_process.html#making-process-spanning-arrays-from-external-data).
If a variable has a particular layout (e.g. is sharded across multiple devices/hosts), then any corresponding auxiliary variables, like optimizer gradient accumulators, need to have the same layout. This requires use to set the variable layout prior to initialization so that it is initialized correctly and efficiently across devices. Added optional `kwargs` to the base `Variable` class so they can handle (and ignore) any backend-specific options. Modified `JaxVariable` to allow setting the layout parameter on construction. Modified `add_variable_from_reference` to copy the layout from the reference variable.
|
Thanks for the PR! LGTM Do you know why tests are failing on CI? Probably some merge issue?
|
fchollet
approved these changes
May 12, 2025
google-ml-butler
bot
added
kokoro:force-run
ready to pull
|
@fchollet yeah, I think it's something do with the api_gen, trying to sort it out |
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.
This continues #19571
This PR includes the following:
angle,bartlett, andblackmanLast local test run (minus applications)
10990 passed, 2022 skipped, 3 xfailed, 7 xpassed in 186.60s (0:03:06)