fix(exla): support container arguments in shard_jit input_shardings - #1834
Merged
Merged
Conversation
polvalente
marked this pull request as draft
September 7, 2026 16:22
polvalente
marked this pull request as ready for review
September 7, 2026 19:11
… buffers EXLA.MLIR.Module.compile/5 sets executable.device_id to -1 for sharded (SPMD) executables, since they are not pinned to one device. maybe_outfeed passed that -1 straight into EXLA.Defn.Buffers.from_nx!/2, which compares it against each input buffer's device id and, on a mismatch, calls EXLA.DeviceBuffer.copy_to_device/3 with device_id -1. The NIF rejects -1, raising "No matching device found for device_id -1" any time a shard_jit input tensor already lived on an EXLA device (e.g. default_backend set to EXLA.Backend) instead of Nx.BinaryBackend. from_nx!/4 now takes the real target device id for the buffer it is preparing. maybe_outfeed passes each partition's index as that id when the run is sharded, matching PjRt's default per-partition device assignment, and keeps using executable.device_id otherwise. Added a regression test in sharding_test.exs that reproduces the crash with EXLA.Backend-resident inputs and checks the fix returns correct per-partition results with the right device ids.
Replaces the two optional positional args (transfer?, target_device_id) with a single opts list, so call sites read as flags instead of bare booleans/nils.
convert_and_validate_input_shardings!/3 read var.shape directly, so any argument that was a container (an Axon.ModelState, or any other struct/map implementing Nx.Container) instead of a plain Nx.Tensor crashed with an unrelated KeyError. Each sharding spec now applies to every tensor leaf inside the matching argument, found by traversing the argument the same way the rest of the sharding pipeline already does. calculate_unsharded_inputs/2 is updated to match: it now consumes one shape multiplier per leaf instead of one per top-level argument.
…Nx.Defn.Compiler The container-flattening and axis-validation logic added for shard_jit's input_shardings lived entirely in EXLA.Defn. Move it into Nx.Defn.Compiler (validate_and_convert_input_shardings!/3 and calculate_unsharded_inputs/2) so any compiler backend implementing __shard_jit__ gets the same container support and validation for free, instead of having to reimplement it. EXLA.Defn now just delegates to these shared functions.
…ompiler dispatcher Instead of exposing validate_and_convert_input_shardings!/3 and calculate_unsharded_inputs/2 as helpers backends must call themselves, run them once inside Nx.Defn.Compiler.__shard_jit__/5, the shared dispatcher every __shard_jit__ call already goes through. Backends now receive vars with their shapes already adjusted and options[:input_shardings] already validated and expanded to one entry per tensor leaf, the same way to_lazy_params_sharded/2 already flattens args_list for them. EXLA.Defn.__shard_jit__/6 no longer needs to know anything about containers at all: it just merges the mesh into options and compiles.
Container flattening and input_shardings validation for shard_jit now live in Nx.Defn.Compiler, shared by every backend, so they should be tested at the Nx level too rather than only through EXLA. Adds a minimal recording Nx.Defn.Compiler that captures what Nx.Defn.Compiler.__shard_jit__/5 passes down, and asserts a container argument's leaves each get their sharding spec, its shape gets scaled correctly, and invalid specs raise a clear ArgumentError.
polvalente
force-pushed
the
fix/exla-shard-jit-container-input-shardings
branch
from
September 10, 2026 20:24
9cf22dc to
c9461ce
Compare
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Stacked on #1833.
Problem
EXLA.shard_jit/3 read var.shape directly in convert_and_validate_input_shardings!/3 to validate each input_shardings entry against its matching argument. When an argument is a container instead of a plain Nx.Tensor (for example an Axon.ModelState, or any struct/map implementing Nx.Container), this crashed with an unrelated KeyError instead of validating the sharding spec.
EXLA.jit/2 and Nx.Defn.jit/2 already accept containers transparently; shard_jit did not.
Fix
Container flattening now happens once, in Nx.Defn.Compiler.shard_jit/5, the shared dispatcher every backend's shard_jit callback goes through, the same way to_lazy_params_sharded/2 already flattens args_list for regular compilation. Each input_shardings entry is validated and expanded to one entry per tensor leaf, and the argument's shape is adjusted to its unsharded (full logical) size, before a backend ever sees it.
A single %{} (fully replicated) sharding on a container argument applies to every leaf inside it, matching the behavior described in the bug report.
Compiler backends no longer need to know anything about containers. EXLA.Defn.shard_jit/6 shrinks down to merging the mesh into options and compiling; any other compiler implementing shard_jit gets the same container support automatically.
Testing
Added tests in sharding_test.exs covering a map container argument with a replicated sharding spec, and a clear ArgumentError (not a KeyError) when a container leaf's sharding spec is invalid.
Full mix test suite passes for both nx and exla.