Since calculating flops is non-trivial, we don't test all the cases currently tested by `test_conv_general_dilated_unfused_hbm_bytes`. Instead, we test behaviors more directly.
PiperOrigin-RevId: 743272840
These args are optional, so not specifying them in our tests will make them simpler and easier to read. This change is a no-op.
PiperOrigin-RevId: 740015584
We'll need to use batch/feature when calculating flops, so it'll help reduce the size of the "calculating-flops" change if we can include them in our tests now.
PiperOrigin-RevId: 739081930
This subsumes (and ultimately will deprecate) overriding the number of CPU devices via XLA_FLAGS.
In addition, replace the test utility jtu.set_host_platform_device_count with jtu.request_cpu_devices(...), which sets or increases the flag's value. This both removes the need for an overly complicated context stack, and prepares for removing remaining uses of setUpModule as part of work parallelizing the test suite with threads.
PiperOrigin-RevId: 713272197