Some tests fail on A100 (that do not fail on A6000)

[stefan-apollo]

On main (d77a4f830d9d031b8443ccd43369ab68989e4071).

A6000:

collected 236 items                                                                                                                          

tests/test_ablations.py ...............                                                                                                [  6%]
tests/test_build_graph.py ..............................................................................                               [ 39%]
tests/test_data.py ..                                                                                                                  [ 40%]
tests/test_data_accumulator.py ...                                                                                                     [ 41%]
tests/test_distributed.py ssss                                                                                                         [ 43%]
tests/test_float_precision.py .s....                                                                                                         [ 45%]
tests/test_folded_bias.py ....X                                                                                                              [ 47%]
tests/test_hooks.py ....                                                                                                                     [ 49%]
tests/test_linalg.py ...............xx...s............                                                                                       [ 63%]
tests/test_loader.py ...............                                                                                                         [ 69%]
tests/test_mlp.py ..................                                                                                                         [ 77%]
tests/test_model_utils.py ...............                                                                                                    [ 83%]
tests/test_tlens_converter.py ....                                                                                                           [ 85%]
tests/test_train_mnist.py .                                                                                                                  [ 86%]
tests/test_train_modular_arithmetic.py .                                                                                                     [ 86%]
tests/test_utils.py ................................                                                                                         [100%]

A100:

collected 236 items                                                                                                                                                                                  

tests/test_ablations.py ...............                                                                                                                                                        [  6%]
tests/test_build_graph.py FFFFFF........................................................................                                                                                                                                                                                                                                                                                                              [ 39%]
tests/test_data.py ..                                                                                                                                                                                                                                                                                                                                                                                                 [ 40%]
tests/test_data_accumulator.py ...                                                                                                                                                                                                                                                                                                                                                                                    [ 41%]
tests/test_distributed.py ssss                                                                                                                                                                                                                                                                                                                                                                                        [ 43%]
tests/test_float_precision.py .s....                                                                                                                                                                                                                                                                                                                                                                                  [ 45%]
tests/test_folded_bias.py ....X                                                                                                                                                                                                                                                                                                                                                                                       [ 47%]
tests/test_hooks.py ....                                                                                                                                                                                                                                                                                                                                                                                              [ 49%]
tests/test_linalg.py ...............xx...s............                                                                                                                                                                                                                                                                                                                                                                [ 63%]
tests/test_loader.py ...............                                                                                                                                                                                                                                                                                                                                                                                  [ 69%]
tests/test_mlp.py ..................                                                                                                                                                                                                                                                                                                                                                                                  [ 77%]
tests/test_model_utils.py ...............                                                                                                                                                                                                                                                                                                                                                                             [ 83%]
tests/test_tlens_converter.py ....                                                                                                                                                                                                                                                                                                                                                                                    [ 85%]
tests/test_train_mnist.py .                                                                                                                                                                                                                                                                                                                                                                                           [ 86%]
tests/test_train_modular_arithmetic.py .                                                                                                                                                                                                                                                                                                                                                                              [ 86%]
tests/test_utils.py ................................    

Failed tests:

FAILED tests/test_build_graph.py::test_modular_arithmetic_build_graph[(1-alpha)^2-functional-trapezoidal] - AssertionError: Tensor-likes are not close!
FAILED tests/test_build_graph.py::test_modular_arithmetic_build_graph[(1-0)*alpha-functional-trapezoidal] - AssertionError: Tensor-likes are not close!
FAILED tests/test_build_graph.py::test_modular_arithmetic_build_graph[(1-alpha)^2-squared-trapezoidal] - AssertionError: Tensor-likes are not close!
FAILED tests/test_build_graph.py::test_modular_arithmetic_build_graph[(1-0)*alpha-squared-trapezoidal] - AssertionError: Tensor-likes are not close!
FAILED tests/test_build_graph.py::test_modular_arithmetic_build_graph[jacobian-squared-trapezoidal] - AssertionError: Tensor-likes are not close!
FAILED tests/test_build_graph.py::test_modular_arithmetic_build_graph[jacobian-squared-gauss-legendre] - AssertionError: Tensor-likes are not close!