Handle native backend tensors in data adapters - #23596
Conversation
Add ops.is_tensor(x) fallback to can_slice_array and convert_to_sliceable so pluggable backend tensors (e.g. MLX) are accepted by the array data adapter. Uses np.asarray for zero-copy views via the buffer protocol, falling back to ops.convert_to_numpy for dtypes numpy can't handle directly (e.g. MLX bfloat16).
There was a problem hiding this comment.
Code Review
This pull request introduces support for slicing pluggable backend tensors (such as MLX) by checking if the input is a tensor and converting it to a NumPy array. However, the current implementation uses ops.is_tensor(x), which will raise an AttributeError because is_tensor is not exported by keras.src.ops. It is recommended to use backend.is_tensor(x) instead, which is already imported and available.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## pluggable_backend #23596 +/- ##
====================================================
Coverage ? 83.98%
====================================================
Files ? 458
Lines ? 62472
Branches ? 10609
====================================================
Hits ? 52469
Misses ? 7329
Partials ? 2674
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
hertschuh
left a comment
There was a problem hiding this comment.
Thanks for working on this!
CC @MarcosAsh
|
i added three tests. one checks that adapters without a native iterator fall back to the numpy iterator, one checks that native arrays remain native across regular and shuffled batching, and one checks the conversion methods on nativearraysliceable. the small fake array is there so the generic pluggable backend path can be tested in keras without depending on mlx or another backend package. once this goes through, the tests on keras-team/keras-mlx#57 should turn green too. also added the |
hertschuh
left a comment
There was a problem hiding this comment.
Did you get a chance to test that it works with the MLX backend?
|
removed the mock tests and made the torch related changes. previously for the NativeArraySliceable I was using since torch tensors will now use i hope this approach makes sense. |
|
Ran it against keras-mlx#57 on Apple silicon with the CI config, mlx 0.31.2: full keras suite 12938 passed, 0 failed. Trainers and callbacks also pass with compile on. Works for MLX. |
as suggested by @hertschuh in keras-team/keras-mlx#40
Description
the data adapters should handle native arrays where
is_tensor(x)is true, rather than relying on a keras-mlx workaround that converts the full dataset to numpy.Pluggable backend tensors (e.g. MLX arrays) are rejected by the array data adapter.
can_slice_array()andconvert_to_sliceable()check for numpy, TF, JAX, torch, pandas, scipy, and__array__, but MLX intentionally dropped__array__in favor of the Python Buffer Protocol (ml-explore/mlx#323). So MLX arrays fall through every check andmodel.fit/evaluate/predictcrashes.Fix
added
ops.is_tensor(x)as a fallback in bothcan_slice_array()andconvert_to_sliceable(). For the sliceable representation, usingnp.asarray(x)to get a zero-copy numpy view. falling back toops.convert_to_numpy(x)whennp.asarrayraises (e.g. MLX bfloat16, which fails the buffer protocol) or silently wraps the input in a 0-D object array (e.g. openvino tensors, which don't expose__array__or the buffer protocol).Contributor Agreement
Please review our AI-Assisted Contribution Policy and check all boxes below before submitting your PR for review:
Note: Failing to adhere to this agreement may result in your future PRs no longer being reviewed.