Mohsin Mehmood
Mohsin Mehmood
> Your last push undid the changes I made to your branch to fix osx. Yes, I was trying something. I am working on it. Will push the updates soon....
> > > Your last push undid the changes I made to your branch to fix osx. > > > > > > Yes, I was trying something. I am...
> Ok, we are getting there. Newer versions fail because it seems we might not be handling properly the PredefinedExpr expression kind. Perhaps you can build locally a newer version...
I tried to reproduce it, transforms must see Modules as explicit args so split/merge and update_context can track mutations. Passing a bound method (e.g., self.block) captures the Module by closure,...
@samanklesaria, I will open the PR I am just thinking a bit more of edge cases and see if there can be a better solution. Feel free to suggest if...
@samanklesaria, can you please review the PR. Let me know if you have any feedback. I will work on it.
I was able to reproduce the issue, and I think the problem comes from how sharding metadata is handled when NNX transforms add new axes. I think vmap, scan, and...
@samanklesaria Good point on the metadata gap. I’ve put together a fix that makes vmap/scan/pmap insert a default {PARTITION_NAME: None} so sharding metadata tracks the extra axis automatically, and made...
@samanklesaria, cool. let me know after you will discuss with @cgarciae then we can discuss and work on it.
@vfdev-5 I'm not sure if this will be good, but I’m thinking what if we make cached_partial tolerant to plain JAX arrays but treat them as read‑only: only nnx.Variables get...