flax

c4725328 - Start exporting event files to benchmark_output_dir as well, as this will help with easier integration with tensorboard.dev

Start exporting event files to benchmark_output_dir as well, as this will help with easier integration with tensorboard.dev Example: Added a dummy test method in mnist_benchmark, results: https://tensorboard.dev/experiment/vaimAvagQkSg5ZkIxiZ0cg/#scalars Previously, summaries by both the test methods in the benchmark class would be combined into one. PiperOrigin-RevId: 321391614