approaches to handling the different "types" of jax backend non_trainable_variables

[matpalm]

one thing i've been doing is using the jax backend with vmap over a model's stateless_call to compose small models into bigger models with different batching sizes

a nuance of this is that as a return we get a vectorised non_trainable_variables which, depending on the use case of course, probably needs to be reduced to a single set of variables for the next call...

so depending on the effective "type" of the non_trainable_variable this reduction might be done in different ways; e.g. i find myself commonly reducing batch norm stats with code like ...

vmapped_fn = vmap(model.stateless_call, in_axes=[None, None, 0])
y_pred, nt_params = vmapped_fn(params, nt_params, x)
nt_params = jax.tree.map(lambda p: jnp.mean(p, axis=0), nt_params)

( though it's kinda clunky since it ends up being a mean of a mean kinda situation )

i can't do this for all the non_trainable_variables though if i have something like dropout, since it''ll include things like the older style uint32 RNG seeds...

i think i can infer it (safely?) by checking the KerasVariable.path but that seems dangerous...

am i missing something around typing for this variables beyond the path from the model itself?

[fchollet]

The variables in model.non_trainable_variables are KerasVariable instances which have the property aggregation. It can take the values "mean", "sum", and "only_first_replica". It defaults to "mean", so I assume RNG seeds would have that (but we can easily change it if that's a bug -- could be "None" I suppose). It doesn't really matter how you "reduce" RNG seeds as long as each replica sees a new seed at each step.