data_format in JAX/FLAX

45 Views Asked by ujjwalnur At 20 October 2024 at 05:53

I did not find any settings for data_format=channels_first or data_format=channels_last in FLAX modules ( which are based on JAX ).

On the contrary, TensorFlow does have that designation. Does the choice of data_format is irrelevant to JAX performance ?

Unfortunately, I did not find any kind of information on this subject.

There are 1 best solutions below

jakevdp On 06 June 2023 at 11:23 BEST ANSWER

JAX/Flax has no equivalent to the concept of data_format as used in Tensorflow/Keras.