flax
Add data_transform argument to nn.scan and preserve PartitionSpec attributes.
#1411
Merged
Merged