[blasphemetheus]

Makes Nx.block raise when the function given to it reads a tensor from the enclosing scope, instead of producing wrong results.

A block body is traced once and then compiled as a function of the block's arguments. A tensor the body reads from outside, such as a weight from the surrounding defn, is not one of those arguments, so no compiler can resolve it. Today nothing catches this at trace time, and the failure shows up later in different ways:

• The Evaluator returns an unevaluated expression from the forward pass.
• EXLA fails while lowering the block with a missing-parameter error.
• The gradient of the captured tensor comes out as all zeros, with no error.

For example, with a block whose body computes Nx.sin(Nx.dot(x, w)) while only x is passed as an argument, the gradient with respect to w is zeros where the plain function gives nonzero values.

The cause is that block parameters were created with the caller's context, so an operation mixing a parameter with an outer tensor looked like ordinary code. while avoids this by giving its body a context of its own. The fix does the same for blocks, so the existing "expressions come from different contexts" check fires when the block is built. A body that returns an outer tensor unchanged does not go through that check, so it is checked separately. Outside a trace the arguments have no context, and the check is skipped there as while does.

No block that exists today reads tensors from the enclosing scope. The Nx.LinAlg blocks, Take, TopK, the FFT blocks and the custom-call blocks all pass their tensors as arguments and close over options only.

Three tests cover an operation that mixes a parameter with an outer tensor, a body that returns an outer tensor, and the same mistake inside grad. All three fail on main.

Found while reviewing #1865, where a checkpoint body that reads a weight from the enclosing scope hits exactly this. This draft is based on main at b532ba1e.

🤖 Generated with Claude Code

[blasphemetheus]

I don't see a downside to this change, but this is messing with conventions I'm not super familiar with so there might be a reason to not do this.

[polvalente]

I think this is an important fix!

For Nx.block we need to keep in mind that the default implementation might never run, so it shouldn't rely on closures for anything. If a runtime value from outside its scope is needed, it must be passed as an argument