Neil Girdhar

> The only proposed change that I am aware of is changing the signatures of module.init and module.apply to not return FrozenDicts. Sorry, I'm not actually discussing the topic of...

@avital Fair enough. I need to learn Flax better before I can really suggest something. A couple other options: An `at` operator that does this under the covers, so that...

> Nice! Thanks as always, Neil. > > I have suggestions which fall into two main categories: > > 1. whitespace Of course, I'm happy to fix all whitespace to...

> I'm not so sure about this in general. Why not promise even less and say it's just an Iterable? We want to promise the right amount, Yes, you're absolutely...

Done, and thank you, you all are a pleasure to work with as well!

> As one example, jraph does addition on the leaves outputs of tree_flatten calls. But that doesn't typecheck against Sequence, since you can't add Sequences together! That makes sense. That's...

@hawkinsp If you like, I can make a corresponding pull request to tensorflow's jaxlib to make the interface promise lists?

This should pass when Tensorflow [2.10](https://github.com/tensorflow/tensorflow/releases/tag/v2.10.0)'s version of jaxlib is used thanks to https://github.com/tensorflow/tensorflow/pull/54330 being compiled into it.

@mattjj Would it be possible to get this merged now that the tensorflow type annotations are in jaxlib?

Just curious, and didn't know where to ask. Is there a reason that `k.split()` does something completely different than `split(k)` if `k` is a new-style `KeyArray`? Could this be a...