Collectives™ on Stack Overflow
Find centralized, trusted content and collaborate around the technologies you use most.
Learn more about Collectives
0votes
0answers
101views
Overhead of instantiating a flax model
Is it expensive to keep recreating a Flax network, such asclass QNetwork(nn.Module): dim: int @nn.compact def __call__(self, x): x = nn.Dense(120)(x) x = nn.relu(x) ...
- 8,122
1vote
1answer
99views
How to control hyperparameter within flax.nnx loss function using an optax.schedule?
from jax import numpy as jnpfrom jax import randomfrom flax import nnximport optaxfrom matplotlib import pyplot as pltif __name__ == '__main__': shape = (2,55,1) epochs = 123 rngs = ...
0votes
1answer
99views
JAX shard on GPU and shard on CPU in subroutine, all with JIT
Duplicating my question here: https://github.com/google/flax/discussions/4825I want to have a JAX or NNX jitted function that consumes and returns GPU-sharded tensors. However, inside the function, I ...
- 820
1vote
2answers
268views
Does vmap correctly split the RNG keys?
In the following code, when I remove the vmap, I have the right randomized behavior. However, with vmap, I don't anymore. Isn't this supposed to be one of the features of nnx.vmap?import jaximport ...
- 8,086
0votes
1answer
54views
Efficient multi-host TPU dataset processing
I want to train LLM on TPUv4-32 using JAX/Flax. The dataset is stored in a mounted google storage bucket. The dataset (Red-Pajama-v2) consists of 5000 shards, which are stored in .json.gz files: ~/...
0votes
1answer
89views
Differentiable weight setting in flax NNX
I'm doing some experiments with Flax NNX (not Linen!).What I'm trying to do is compute the weights of a network using another network:A hypernetwork receives some input parameters W and outputs a ...
1vote
0answers
75views
LAPACK Inconsistent across multiple different operating systems and devices
DescriptionI have a deterministic program that uses jax, and is heavy on linear algebra operations.I ran this code on CPU, using three different CPUs. Two MacOs Systems (one on Sequoia (M1 Pro), ...
2votes
1answer
108views
How to type hint `flax.linen.Module.apply`'s output correctly?
As of writing, this code does not pass the PyRight type checker:import jaximport jax.numpy as jnpimport jax.typing as jtimport flax.linen as nnclass MLP(nn.Module): @nn.compact def ...
- 431
1vote
1answer
290views
Would using lists rather than jax.numpy arrays lead to more accurate numerical transformations?
I am doing a project with RNNs using jax and flax and I have noticed some behavior that I do not really understand.My code is basically an optimization loop where the user provides the initial ...
1vote
1answer
279views
Is there a way to update weights of an nnx.Module in Flax's NNX using the lax.scan function?
I have a neural network (nnx.Module) written in Flax's NNX. I want to train this network efficiently using lax.scan instead of a for loop. However, as scan doesn't allow in place changes, how can I ...
- 21
1vote
1answer
372views
Freezing filtered parameter collections with Flax.nnx
I'm trying to work out how to do transfer learning with flax.nnx. Below is my attempt to freeze the kernel of my nnx.Linear instance and optimize the bias. I think maybe I'm not correctly setting up ...
- 23
2votes
1answer
203views
Debug jax In vscode
Why can't I use a vscode debugger to debug jax code, specifically pure functions. I understand that they provide their own framework for debugging but vscode debugger is quite comfortable. Is this ...
- 33
0votes
1answer
286views
Flax nnx / jax: tree.map for layers of incongruent size
I am trying to figure out how to use nnx.split_rngs. Can somebody give a version of the code below that uses nnx.split_rngs with jax.tree.map to produce an arbitrary number of Linear layers with ...
- 23
1vote
1answer
90views
Jax / Flax potential tracing issue
I'm currently using Flax for neural network implementations. My model takes two inputs:x and θ. It first processes x through an LSTM, then concatenates the LSTM's output with θ — or more precisely, ...
0votes
0answers
76views
orbax save/restore with 8 devices
I have setup a snippet on Colab herewithjax.__version__ # 0.4.33 9Feb2025 orbax.checkpoint.__version__ # 0.6.4 9Feb2025It quite difficult to follow the flax/orbax changes in the ...
- 402
- The Overflow Blog
- Featured on Meta
Hot Network Questions
- How to find parent ways of nodes
- Not Benny but close
- Problem with the proof of Theorem 5.1.2 in Mac Lane's "Categories for the Working Mathematician"?
- Why do my Raspberry Pi 5 SD cards have unwritable partition tables?
- Is a proper wire nut connection possible if wires don’t twist?
- Manga involving the MC playing a computer game connected to a medieval fantasy world
- What is the maximum number of damage types for a single attack?
- Is there a term for the sound the orchestra produces when they tune?
- Difference between VOR approach (with DME required) and VOR DME approach?
- When was the symbol "U" first used to denote voltage or electric potential difference?
- What physical changes to the mouth would cause problems with articulation?
- What are the events in Parliament Square and Sieges Allee that Woolf refers to in A Room Of One's Own?
- Number of subgroups in a semidirect product
- A novel where a time traveler arrived by mistake at the peak of the Black Death
- What does "origin of the world" vulgo "Big Bang" mean?
- Diagnosing non-operational wired doorbell
- Low-Side Current Sense Simulation
- Detect if any ping succeeds
- Why does JavaScript use autoboxing?
- Is there an asymmetric standard applied to zero probability events?
- What damages are recoverable for late delivery of a package?
- Does the intersection of a variety with rational singularities with a smooth variety have rational singularities?
- Will this this schematic modification work?
- Was the spittle and mud in John 9:6-7, symbolic, a catalyst to faith, or a required act of obedience?